diff --git a/gaussian_renderer/__init__.py b/gaussian_renderer/__init__.py index e12f4b6803..4f65dc6775 100644 --- a/gaussian_renderer/__init__.py +++ b/gaussian_renderer/__init__.py @@ -126,3 +126,71 @@ def render(viewpoint_camera, pc : GaussianModel, pipe, bg_color : torch.Tensor, } return out + + + + +##################### 추가함수 ##################### +def render_reflected_gaussians(viewpoint_camera, pc: GaussianModel, pipe, bg_color: torch.Tensor, mirror_transform: torch.Tensor, gaussians_to_reflect_mask: torch.Tensor, scaling_modifier=1.0): + reflected_attrs = pc.reflect(mirror_transform, gaussians_to_reflect_mask) + + if reflected_attrs is None or reflected_attrs["xyz"].shape[0] == 0: + image_height = int(viewpoint_camera.image_height) + image_width = int(viewpoint_camera.image_width) + return {"render": torch.full((3, image_height, image_width), bg_color[0].item(), device="cuda"), "radii": torch.zeros(0, device="cuda")} + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + # 카메라 시선벡터 반사 + orig_campos = viewpoint_camera.camera_center + campos_hom = torch.cat([orig_campos, torch.ones(1, device="cuda")], dim=0) + reflected_campos_hom = campos_hom @ mirror_transform.T + reflected_campos = reflected_campos_hom[:3] + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=reflected_campos, + prefiltered=False, + debug=pipe.debug, + antialiasing=pipe.antialiasing + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = reflected_attrs["xyz"] + rotations = torch.nn.functional.normalize(reflected_attrs["rotation"]) + scales = torch.exp(reflected_attrs["scaling"]) + opacity = torch.sigmoid(reflected_attrs["opacity"]) + shs = torch.cat((reflected_attrs["features_dc"], reflected_attrs["features_rest"]), dim=1) + + screenspace_points = torch.zeros_like(means3D, dtype=means3D.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + # Rasterize + rendered_image, radii, _ = rasterizer( + means3D = means3D, + means2D = screenspace_points, + shs = shs, + colors_precomp = None, + opacities = opacity, + scales = scales, + rotations = rotations, + cov3D_precomp = None + ) + + return { + "render": rendered_image.clamp(0, 1), + "radii": radii + } diff --git a/propose_mirror_plane.py b/propose_mirror_plane.py new file mode 100644 index 0000000000..aae0798159 --- /dev/null +++ b/propose_mirror_plane.py @@ -0,0 +1,57 @@ +import numpy as np +import os +import json +from argparse import ArgumentParser +import sys + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '.'))) + +from scene.colmap_loader import read_points3D_binary + +def propose_plane_robust(points_bin_path: str, output_json_path: str): + print(f"Analyzing point cloud from: {points_bin_path}") + try: + xyz, _, _ = read_points3D_binary(points_bin_path) + if xyz.shape[0] == 0: + raise ValueError("Point cloud is empty.") + + y_coords = xyz[:, 1] + + # 모든 가우시안의 y값의 하위 5%를 실질적 바닥으로 둠 + y_percentile_5 = np.percentile(y_coords, 5) + y_mean = np.mean(y_coords) + + print(f"Point cloud Y-axis 5th percentile: {y_percentile_5:.3f}") + print(f"Point cloud Y-axis mean: {y_mean:.3f}") + + # 바닥에서 y값 평균 ~ 바닥 거리 만큼 뺌 -> 바닥 아래에 반사 평면 위치시킴 + plane_y_value = y_percentile_5 - (y_mean - y_percentile_5) + + # 법선벡터는 (0, 1, 0)으로 고정 + plane_params = { + "a": 0.0, + "b": 1.0, + "c": 0.0, + "d": -plane_y_value + } + + os.makedirs(os.path.dirname(output_json_path), exist_ok=True) + with open(output_json_path, 'w') as f: + json.dump(plane_params, f, indent=4) + + print(f"\nSuccessfully proposed robust mirror plane: y = {plane_y_value:.3f}") + print(f"Plane parameters saved to: {output_json_path}") + + except Exception as e: + print(f"Error: Failed to propose mirror plane. {e}") + +if __name__ == "__main__": + parser = ArgumentParser(description="Propose a robust mirror plane from a COLMAP sparse point cloud.") + parser.add_argument("-s", "--source_path", required=True, type=str, help="Path to the COLMAP dataset directory") + parser.add_argument("-m", "--model_path", required=True, type=str, help="Path to the output model directory where the plane file will be saved") + args = parser.parse_args() + + points_3d_bin = os.path.join(args.source_path, "sparse/0/points3D.bin") + output_json = os.path.join(args.model_path, "mirror_plane.json") + + propose_plane_robust(points_3d_bin, output_json) diff --git a/scene/gaussian_model.py b/scene/gaussian_model.py index 473887db89..9b73e0d902 100644 --- a/scene/gaussian_model.py +++ b/scene/gaussian_model.py @@ -19,8 +19,8 @@ from plyfile import PlyData, PlyElement from utils.sh_utils import RGB2SH from simple_knn._C import distCUDA2 -from utils.graphics_utils import BasicPointCloud -from utils.general_utils import strip_symmetric, build_scaling_rotation +from utils.graphics_utils import BasicPointCloud, geom_transform_points +from utils.general_utils import strip_symmetric, build_scaling_rotation, quat_from_matrix, build_rotation try: from diff_gaussian_rasterization import SparseGaussianAdam @@ -318,14 +318,18 @@ def replace_tensor_to_optimizer(self, tensor, name): for group in self.optimizer.param_groups: if group["name"] == name: stored_state = self.optimizer.state.get(group['params'][0], None) - stored_state["exp_avg"] = torch.zeros_like(tensor) - stored_state["exp_avg_sq"] = torch.zeros_like(tensor) - - del self.optimizer.state[group['params'][0]] - group["params"][0] = nn.Parameter(tensor.requires_grad_(True)) - self.optimizer.state[group['params'][0]] = stored_state - - optimizable_tensors[group["name"]] = group["params"][0] + if stored_state is not None: + stored_state["exp_avg"] = torch.zeros_like(tensor) + stored_state["exp_avg_sq"] = torch.zeros_like(tensor) + + del self.optimizer.state[group['params'][0]] + group["params"][0] = nn.Parameter(tensor.requires_grad_(True)) + self.optimizer.state[group['params'][0]] = stored_state + + optimizable_tensors[group["name"]] = group["params"][0] + else: + group["params"][0] = nn.Parameter(tensor.requires_grad_(True)) + optimizable_tensors[group["name"]] = group["params"][0] return optimizable_tensors def _prune_optimizer(self, mask): @@ -471,3 +475,46 @@ def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size, radi def add_densification_stats(self, viewspace_point_tensor, update_filter): self.xyz_gradient_accum[update_filter] += torch.norm(viewspace_point_tensor.grad[update_filter,:2], dim=-1, keepdim=True) self.denom[update_filter] += 1 + + + + +##################### 추가함수 ##################### + @torch.no_grad() + def reflect(self, mirror_transform, reflect_mask): + from utils.graphics_utils import geom_transform_points + from utils.general_utils import build_rotation + + if reflect_mask.sum() == 0: + return None + + object_rotation_q = self.get_rotation[reflect_mask] + object_xyz = self.get_xyz[reflect_mask] + + # 센터 반사 + reflected_xyz = geom_transform_points(object_xyz, mirror_transform.T) + + # 회전 반사 + H_3x3 = mirror_transform[:3, :3] + R = build_rotation(object_rotation_q) + R_reflected = H_3x3 @ R + # 재직교화(SVD) + det(+1) 강제 + U, S, Vt = torch.linalg.svd(R_reflected) # (N,3,3) + R_reflected = U @ Vt + det = torch.det(R_reflected) + neg = det < 0 + if neg.any(): + Vt[neg, -1, :] *= -1.0 + R_reflected = U @ Vt + + reflected_rotation_q = quat_from_matrix(R_reflected) + + # 반사시킨 속성 반환 + return { + "xyz": reflected_xyz, + "rotation": torch.nn.functional.normalize(reflected_rotation_q, p=2, dim=1), + "scaling": self._scaling[reflect_mask], + "opacity": self._opacity[reflect_mask], + "features_dc": self._features_dc[reflect_mask], + "features_rest": self._features_rest[reflect_mask] + } diff --git a/toy_render.py b/toy_render.py new file mode 100644 index 0000000000..9b05a65cac --- /dev/null +++ b/toy_render.py @@ -0,0 +1,193 @@ +import os +import sys +import torch +import numpy as np +import json +from argparse import ArgumentParser +from tqdm import tqdm +import torchvision +import plyfile + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '.'))) + +from scene import Scene, GaussianModel +from gaussian_renderer import render, render_reflected_gaussians +from utils.general_utils import safe_state +from arguments import ModelParams, PipelineParams, get_combined_args + + +# PLY 파일로 저장 +def save_ply_from_attrs(path, attrs): + os.makedirs(os.path.dirname(path), exist_ok=True) + + xyz = attrs["xyz"] + + + xyz_np = xyz.detach().cpu().numpy() + f_dc_np = attrs["features_dc"].detach().cpu().numpy().reshape(-1, 3) + f_rest_np = attrs["features_rest"].detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() + opacities_np = torch.sigmoid(attrs["opacity"]).detach().cpu().numpy() + scales_np = torch.exp(attrs["scaling"]).detach().cpu().numpy() + rotations_np = torch.nn.functional.normalize(attrs["rotation"]).detach().cpu().numpy() + + + dtype_full = [ + ('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('nx', 'f4'), ('ny', 'f4'), ('nz', 'f4'), + ('f_dc_0', 'f4'), ('f_dc_1', 'f4'), ('f_dc_2', 'f4') + ] + for i in range(f_rest_np.shape[1]): + dtype_full.append((f'f_rest_{i}', 'f4')) + dtype_full.extend([ + ('opacity', 'f4'), + ('scale_0', 'f4'), ('scale_1', 'f4'), ('scale_2', 'f4'), + ('rot_0', 'f4'), ('rot_1', 'f4'), ('rot_2', 'f4'), ('rot_3', 'f4') + ]) + + normals = np.zeros_like(xyz_np) + elements = np.empty(xyz_np.shape[0], dtype=dtype_full) + attributes = np.concatenate((xyz_np, normals, f_dc_np, f_rest_np, opacities_np, scales_np, rotations_np), axis=1) + elements[:] = list(map(tuple, attributes)) + + el = plyfile.PlyElement.describe(elements, 'vertex') + plyfile.PlyData([el], text=True).write(path) + print(f"Combined gaussians saved to {path}") + +# 원본 + 반사 가우시안 렌더링 +def render_reflection_set(model_path, name, iteration, views, gaussians, pipeline, background, mirror_transform): + render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders") + gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt") + + os.makedirs(render_path, exist_ok=True) + os.makedirs(gts_path, exist_ok=True) + + for idx, view in enumerate(tqdm(views, desc="Rendering progress")): + # 불투명도 > 0.1 인 가우시안만 반사하여 유의미한 객체만 선택 + opacity_threshold = 0.1 + reflect_mask = (gaussians.get_opacity > opacity_threshold).squeeze() + reflected_attrs = gaussians.reflect(mirror_transform, reflect_mask) + + # 원본 가우시안 속성 + original_xyz = gaussians._xyz + original_rotation = gaussians._rotation + original_scaling = gaussians._scaling + original_opacity = gaussians._opacity + original_features_dc = gaussians._features_dc + original_features_rest = gaussians._features_rest + + combined_gaussians = GaussianModel(gaussians.max_sh_degree) + + if reflected_attrs is not None: + # 원본과 반사된 가우시안 속성 결합 + combined_xyz = torch.cat((original_xyz, reflected_attrs["xyz"]), dim=0) + combined_rotation = torch.cat((original_rotation, reflected_attrs["rotation"]), dim=0) + combined_scaling = torch.cat((original_scaling, reflected_attrs["scaling"]), dim=0) + combined_opacity = torch.cat((original_opacity, reflected_attrs["opacity"]), dim=0) + combined_features_dc = torch.cat((original_features_dc, reflected_attrs["features_dc"]), dim=0) + combined_features_rest = torch.cat((original_features_rest, reflected_attrs["features_rest"]), dim=0) + else: + # 반사된 가우시안이 없으면 원본만 사용 + combined_xyz = original_xyz + combined_rotation = original_rotation + combined_scaling = original_scaling + combined_opacity = original_opacity + combined_features_dc = original_features_dc + combined_features_rest = original_features_rest + + # 결합된 가우시안으로 임시 GaussianModel 인스턴스 생성 + combined_gaussians._xyz = torch.nn.Parameter(combined_xyz.detach().requires_grad_(False)) + combined_gaussians._rotation = torch.nn.Parameter(combined_rotation.detach().requires_grad_(False)) + combined_gaussians._scaling = torch.nn.Parameter(combined_scaling.detach().requires_grad_(False)) + combined_gaussians._opacity = torch.nn.Parameter(combined_opacity.detach().requires_grad_(False)) + combined_gaussians._features_dc = torch.nn.Parameter(combined_features_dc.detach().requires_grad_(False)) + combined_gaussians._features_rest = torch.nn.Parameter(combined_features_rest.detach().requires_grad_(False)) + combined_gaussians.active_sh_degree = gaussians.active_sh_degree # SH degree도 원본과 동일하게 설정 + + # 결합된 가우시안을 렌더링 + final_image = render(view, combined_gaussians, pipeline, background)["render"] + gt = view.original_image[0:3, :, :] + + torchvision.utils.save_image(final_image, os.path.join(render_path, '{0:05d}'.format(idx) + ".png")) + torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ".png")) + +# 렌더링 준비 +def render_sets(dataset: ModelParams, iteration: int, pipeline: PipelineParams, skip_train: bool, skip_test: bool): + with torch.no_grad(): + gaussians = GaussianModel(dataset.sh_degree) + scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False) + + bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0] + background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + + plane_json_path = os.path.join(dataset.model_path, "mirror_plane.json") + if not os.path.exists(plane_json_path): + raise FileNotFoundError(f"Mirror plane file not found at {plane_json_path}.") + + with open(plane_json_path, 'r') as f: + plane_params = json.load(f) + + a, b, c, d = plane_params['a'], plane_params['b'], plane_params['c'], plane_params['d'] + + # 평면 정규화 추가 + n = np.array([a, b, c], dtype=np.float32) + s = np.linalg.norm(n) + 1e-12 + n /= s + d /= s + a, b, c = n.tolist() + + # 정규화된 하우스홀더 반사 행렬 + H = np.array([ + [1 - 2*a*a, -2*a*b, -2*a*c, -2*a*d], + [-2*a*b, 1 - 2*b*b, -2*b*c, -2*b*d], + [-2*a*c, -2*b*c, 1 - 2*c*c, -2*c*d], + [0, 0, 0, 1] + ], dtype=np.float32) + + mirror_transform = torch.from_numpy(H).cuda() + + print("\nExporting PLY files...") + original_attrs = { + "xyz": gaussians.get_xyz, + "rotation": gaussians._rotation, + "scaling": gaussians._scaling, + "opacity": gaussians._opacity, + "features_dc": gaussians._features_dc, + "features_rest": gaussians._features_rest + } + # PLY 내보내기를 위해 모든 가우시안을 포함하는 마스크 사용 + reflect_mask_all = torch.ones(gaussians.get_xyz.shape[0], dtype=torch.bool, device="cuda") + reflected_attrs = gaussians.reflect(mirror_transform, reflect_mask_all) + + if reflected_attrs: + # 1. 반사된 가우시안만 저장 + reflected_ply_path = os.path.join(dataset.model_path, f"reflected_gaussians_iter{iteration}.ply") + save_ply_from_attrs(reflected_ply_path, reflected_attrs) + print(f"\n[SUCCESS] Reflected gaussians saved to: {reflected_ply_path}") + + # 2. 원본 + 반사된 가우시안 저장 + combined_attrs = {key: torch.cat((original_attrs[key], reflected_attrs[key]), dim=0) for key in original_attrs} + combined_ply_path = os.path.join(dataset.model_path, f"combined_gaussians_iter{iteration}.ply") + save_ply_from_attrs(combined_ply_path, combined_attrs) + print(f"\n[SUCCESS] Combined gaussians saved to: {combined_ply_path}") + else: + print("[INFO] No gaussians were reflected. PLY export skipped.") + + if not skip_train: + render_reflection_set(dataset.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background, mirror_transform) + + if not skip_test: + render_reflection_set(dataset.model_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, pipeline, background, mirror_transform) + +if __name__ == "__main__": + parser = ArgumentParser(description="Render original and reflected gaussians and save PLY files.") + model = ModelParams(parser, sentinel=True) + pipeline = PipelineParams(parser) + parser.add_argument("--iteration", default=7_000, type=int, help="Iteration number of the model to load.") + parser.add_argument("--skip_train", action="store_true") + parser.add_argument("--skip_test", action="store_true") + parser.add_argument("--quiet", action="store_true") + args = get_combined_args(parser) + + print("Rendering " + args.model_path) + safe_state(args.quiet) + + render_sets(model.extract(args), args.iteration, pipeline.extract(args), args.skip_train, args.skip_test) diff --git a/utils/general_utils.py b/utils/general_utils.py index 541c082522..0007149ff3 100644 --- a/utils/general_utils.py +++ b/utils/general_utils.py @@ -131,3 +131,26 @@ def flush(self): np.random.seed(0) torch.manual_seed(0) torch.cuda.set_device(torch.device("cuda:0")) + +def focal2fov(focal, pixels): + return 2*math.atan(pixels/(2*focal)) + + + + + +##################### 추가함수 ##################### + +def quat_normalize(q): + return q / (q.norm(dim=-1, keepdim=True) + 1e-12) + +def quat_from_matrix(R): # R: (...,3,3) -> (...,4) (w,x,y,z) + t = R[...,0,0] + R[...,1,1] + R[...,2,2] + w = torch.sqrt(torch.clamp(1.0 + t, min=0)) / 2.0 + x = torch.sqrt(torch.clamp(1.0 + R[...,0,0] - R[...,1,1] - R[...,2,2], min=0)) / 2.0 + y = torch.sqrt(torch.clamp(1.0 - R[...,0,0] + R[...,1,1] - R[...,2,2], min=0)) / 2.0 + z = torch.sqrt(torch.clamp(1.0 - R[...,0,0] - R[...,1,1] + R[...,2,2], min=0)) / 2.0 + x = x.copysign(R[...,2,1] - R[...,1,2]) + y = y.copysign(R[...,0,2] - R[...,2,0]) + z = z.copysign(R[...,1,0] - R[...,0,1]) + return quat_normalize(torch.stack([w,x,y,z], dim=-1)) diff --git "a/\354\227\260\354\212\265\354\236\245.py" "b/\354\227\260\354\212\265\354\236\245.py" new file mode 100644 index 0000000000..6befbd1c39 --- /dev/null +++ "b/\354\227\260\354\212\265\354\236\245.py" @@ -0,0 +1,43 @@ +def plus(x): + return x + 1 + + +class Test: + def setup_functions(self): + self.x_plus = plus + self.y_plus = plus + self.z_plus = plus + + def __init__(self): + self.x = 10 + self.y = 20 + self.z = 30 + self.setup_functions() + + def get_plus_x(self): + return self.x_plus(self.x) + + def get_plus_y(self): + return self.y_plus(self.y) + + def get_plus_z(self): + return self.z_plus(self.z) + + +t = Test() + +print(t.x) +print(t.y) +print(t.z) + +print("Changed x: ", t.get_plus_x()) +print("Changed y: ", t.get_plus_y()) +print("Changed z: ", t.get_plus_z()) + +t.x = 100 +t.y = 200 +t.z = 300 + +print("Changed x: ", t.get_plus_x()) +print("Changed y: ", t.get_plus_y()) +print("Changed z: ", t.get_plus_z())