Merge pull request #115 from hustvl/pertimestamp_3dgaussians

add-export-pertimestamp-gaussians

README.md

@@ -162,18 +162,19 @@ python metrics.py --model_path "output/dnerf/bouncingballs/"
 [Watch me](./docs/viewer_usage.md)
 ## Scripts
 
-There are some helpful scripts in , please feel free to use them.
+There are some helpful scripts, please feel free to use them.
 
-`vis_point.py`:
-get all points clouds at each timestamps.
+`export_perframe_3DGS.py`:
+get all 3D Gaussians point clouds at each timestamps.
 
 usage:
 
 ```python
-export exp_name="hypernerf"
-python vis_point.py --model_path output/$exp_name/interp/aleks-teapot --configs arguments/$exp_name/default.py 
+python export_perframe_3DGS.py --iteration 14000 --configs arguments/dnerf/lego.py --model_path output/dnerf/lego 
 ```
 
+You will a set of 3D Gaussians are saved in `output/dnerf/lego/gaussian_pertimestamp`.
+
 `weight_visualization.ipynb`:
 
 visualize the weight of Multi-resolution HexPlane module.

export_perframe_3DGS.py

@@ -14,6 +14,7 @@ from arguments import ModelParams, PipelineParams, get_combined_args, ModelHidde
 from gaussian_renderer import GaussianModel
 from time import time
 import open3d as o3d
+from plyfile import PlyData, PlyElement
 # import torch.multiprocessing as mp
 import threading
 from utils.render_utils import get_state_at_time
@@ -37,6 +38,40 @@ def save_point_cloud(points, model_path, timestamp):
     pcd.points = o3d.utility.Vector3dVector(points)
     ply_path = os.path.join(output_path,f"points_{timestamp}.ply")
     o3d.io.write_point_cloud(ply_path, pcd)
+def construct_list_of_attributes(feature_dc_shape, feature_rest_shape, scaling_shape,rotation_shape):
+    l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
+    # All channels except the 3 DC
+    for i in range(feature_dc_shape[1]*feature_dc_shape[2]):
+        l.append('f_dc_{}'.format(i))
+    for i in range(feature_rest_shape[1]*feature_rest_shape[2]):
+        l.append('f_rest_{}'.format(i))
+    l.append('opacity')
+    for i in range(scaling_shape[1]):
+        l.append('scale_{}'.format(i))
+    for i in range(rotation_shape[1]):
+        l.append('rot_{}'.format(i))
+    # breakpoint()
+    return l
+def init_3DGaussians_ply(points, scales, rotations, opactiy, shs, feature_shape):
+    xyz = points.detach().cpu().numpy()
+    normals = np.zeros_like(xyz)
+    feature_dc = shs[:,0:feature_shape[0],:]
+    feature_rest = shs[:,feature_shape[0]:,:]
+    f_dc = shs[:,:feature_shape[0],:].detach().transpose(1,2).flatten(start_dim=1).contiguous().cpu().numpy()
+    # breakpoint()
+    f_rest = shs[:,feature_shape[0]:,:].detach().transpose(1,2).flatten(start_dim=1).contiguous().cpu().numpy()
+    opacities = opactiy.detach().cpu().numpy()
+    scale = scales.detach().cpu().numpy()
+    rotation = rotations.detach().cpu().numpy()
+
+    dtype_full = [(attribute, 'f4') for attribute in construct_list_of_attributes(feature_dc.shape, feature_rest.shape, scales.shape, rotations.shape)]
+    elements = np.empty(xyz.shape[0], dtype=dtype_full)
+    attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
+    elements[:] = list(map(tuple, attributes))
+    el = PlyElement.describe(elements, 'vertex')
+    # breakpoint()
+    return PlyData([el])
+    
 parser = ArgumentParser(description="Testing script parameters")
 model = ModelParams(parser, sentinel=True)
 pipeline = PipelineParams(parser)
@@ -47,6 +82,8 @@ parser.add_argument("--skip_test", action="store_true")
 parser.add_argument("--quiet", action="store_true")
 parser.add_argument("--skip_video", action="store_true")
 parser.add_argument("--configs", type=str)
+# parser.add_argument("--model_path", type=str)
+
 args = get_combined_args(parser)
 print("Rendering " , args.model_path)
 if args.configs:
@@ -57,6 +94,14 @@ if args.configs:
 # Initialize system state (RNG)
 safe_state(args.quiet)
 gaussians, scene = render_sets(model.extract(args), hyperparam.extract(args), args.iteration, pipeline.extract(args), args.skip_train, args.skip_test, args.skip_video)
-for index, viewpoint in enumerate(scene.getVideoCameras()):
+output_path = os.path.join(args.model_path,"gaussian_pertimestamp")
+os.makedirs(output_path,exist_ok=True)
+print("Computing Gaussians.")
+for index, viewpoint in enumerate(scene.getTestCameras()):
+    
     points, scales_final, rotations_final, opacity_final, shs_final = get_state_at_time(gaussians, viewpoint)
-    save_point_cloud(points, args.model_path, index)
+    feature_dc_shape = gaussians._features_dc.shape[1]
+    feature_rest_shape = gaussians._features_rest.shape[1]
+    gs_ply = init_3DGaussians_ply(points, scales_final, rotations_final, opacity_final, shs_final, [feature_dc_shape, feature_rest_shape])
+    gs_ply.write(os.path.join(output_path,"time_{0:05d}.ply".format(index)))
+print("done")

merge_many_4dgs.py

@@ -105,7 +105,9 @@ def render(viewpoint_camera, gaussians, bg_color : torch.Tensor, scaling_modifie
     for index, pc in enumerate(gaussians):
         
         means3D_final1, scales_final1, rotations_final1, opacity_final1, shs_final1 = get_state_at_time(pc, viewpoint_camera)
-
+        scales_final1 = pc.scaling_activation(scales_final1)
+        rotations_final1 = pc.rotation_activation(rotations_final1)
+        opacity_final1 = pc.opacity_activation(opacity_final1)
         if index == 0:
             means3D_final, scales_final, rotations_final, opacity_final, shs_final = means3D_final1, scales_final1, rotations_final1, opacity_final1, shs_final1
         else:

utils/render_utils.py

@@ -19,7 +19,7 @@ def get_state_at_time(pc,viewpoint_camera):
     means3D_final, scales_final, rotations_final, opacity_final, shs_final = pc._deformation(means3D, scales, 
                                                                  rotations, opacity, shs,
                                                                  time)
-    scales_final = pc.scaling_activation(scales_final)
-    rotations_final = pc.rotation_activation(rotations_final)
-    opacity = pc.opacity_activation(opacity_final)
-    return means3D_final, scales_final, rotations_final, opacity, shs
+    # scales_final = pc.scaling_activation(scales_final)
+    # rotations_final = pc.rotation_activation(rotations_final)
+    # opacity = pc.opacity_activation(opacity_final)
+    return means3D_final, scales_final, rotations_final, opacity, shs_final