misc: enhance simple_trainer with viewer integration and GOF configuration

Author: JasonLSC

examples/datasets/INVR_N3D.py

@@ -24,7 +24,7 @@
     from examples.helper.STG.time_utils import timer, timeblock
 
 # reference to STG's scene __init__.py
-@timer
+# @timer
 class Parser:
     """COLMAP parser."""
 
@@ -56,14 +56,15 @@ def __init__(
         raydict = {}
 
         # Get scene info
-        if loader == "colmap": # colmapvalid only for testing
+        ## Get cam parameters & merged point cloud for splats initialization
+        if loader == "colmap": 
             scene_info = sceneLoadTypeCallbacks["Colmap"](self.source_path, self.images_phrase, self.eval, multiview, duration=self.duration, test_view_id=self.test_view_id, downscale_factor=downscale_factor) # SceneInfo() - NamedTuple
         # elif loader == "invr":
         #     scene_info = sceneLoadTypeCallbacks["INVR"](self.source_path, self.images_phrase, self.eval, multiview, duration=self.duration) # SceneInfo() - NamedTuple
         else:
             assert False, "Could not recognize scene type!"
 
-        with open(scene_info.ply_path, 'rb') as src_file, open(os.path.join(self.model_path, "input.ply") , 'wb') as dest_file:
+        with open(scene_info.ply_path, 'rb') as src_file, open(os.path.join(self.model_path, "init_pcd.ply") , 'wb') as dest_file:
             dest_file.write(src_file.read())
 
         self.cameras_extent = scene_info.nerf_normalization["radius"]
@@ -187,7 +188,8 @@ def __getitem__(self, index: int) -> Dict[str, Any]:
             images.append(PILtoTorch_new(self.fetch_image(finfo.image_path), resolution).permute(1,2,0))
             image_paths.append(finfo.image_path)
             camtoworlds.append(torch.from_numpy(self.camtoworld[globalid]))
-            timesteps.append(tid/len(self.scene_by_t))
+            # timesteps.append(tid/len(self.scene_by_t)) # old version: norm to [0,1]
+            timesteps.append(tid/30) # default framerate is 30 
             Ks.append(torch.from_numpy(self.parser.K))
             rays.append(cami.rays[0])
 

examples/helper/STG/dataset_readers.py

@@ -124,10 +124,6 @@ def readColmapCameras(cam_extrinsics, cam_intrinsics, images_folder, near, far,
         sys.stdout.flush()
 
         extr = cam_extrinsics[key]
-
-
-
-
         intr = cam_intrinsics[extr.camera_id]
         height = intr.height
         width = intr.width
@@ -147,8 +143,6 @@ def readColmapCameras(cam_extrinsics, cam_intrinsics, images_folder, near, far,
         else:
             assert False, "Colmap camera model not handled: only undistorted datasets (PINHOLE or SIMPLE_PINHOLE cameras) supported!"
 
-
-
        # loop for timestamps/frames
         for j in range(startime, startime+ int(duration)):
             image_path = os.path.join(images_folder, os.path.basename(extr.name))
@@ -757,6 +751,7 @@ def readColmapSceneInfoMv(path, images, eval, llffhold=8, multiview=False, durat
 
 
 def readColmapSceneInfo(path, images, eval, llffhold=8, multiview=False, duration=50, test_view_id=[0], downscale_factor=1):
+    # get camera extrinsic & intrinsic
     try:
         cameras_extrinsic_file = os.path.join(path, "sparse/0", "images.bin")
         cameras_intrinsic_file = os.path.join(path, "sparse/0", "cameras.bin")
@@ -787,12 +782,14 @@ def readColmapSceneInfo(path, images, eval, llffhold=8, multiview=False, duratio
             train_cam_infos =  cam_infos[duration:] # Camera 1~20 train set
             test_cam_infos = cam_infos[:duration] # Camera 0 test set
 
+            # make sure only have one test view 
             uniquecheck = []
             for cam_info in test_cam_infos:
                 if cam_info.image_name not in uniquecheck:
                     uniquecheck.append(cam_info.image_name)
             assert len(uniquecheck) == 1 
 
+            # make sure test view is not in the set of train view 
             sanitycheck = []
             for cam_info in train_cam_infos:
                 if cam_info.image_name not in sanitycheck:
@@ -829,12 +826,9 @@ def readColmapSceneInfo(path, images, eval, llffhold=8, multiview=False, duratio
     nerf_normalization = getNerfppNorm(train_cam_infos)
     nerf_normalization_test = getNerfppNorm(test_cam_infos)
 
-    ply_path = os.path.join(path, "sparse/0/points3D.ply")
-    bin_path = os.path.join(path, "sparse/0/points3D.bin")
-    txt_path = os.path.join(path, "sparse/0/points3D.txt")
     totalply_path = os.path.join(path, "sparse/0/points3D_total" + str(duration) + ".ply")
 
-
+    ### operation from STG
     # merge SfM point clouds from consecutive frames into a single .ply file
     # pc from different frames have corresponding time index
     # these points will be used for initialization
@@ -849,14 +843,17 @@ def readColmapSceneInfo(path, images, eval, llffhold=8, multiview=False, duratio
             xyz, rgb, _ = read_points3D_binary(thisbin_path)
             totalxyz.append(xyz)
             totalrgb.append(rgb)
-            totaltime.append(np.ones((xyz.shape[0], 1)) * (i-starttime) / duration)
+            # totaltime.append(np.ones((xyz.shape[0], 1)) * (i-starttime) / duration) # old version: norm timestamp into [0,1]
+            totaltime.append(np.ones((xyz.shape[0], 1)) * (i-starttime) / 30) # new version:change the unit from index into seconds, assume framerate as 30 frames per second
+
         xyz = np.concatenate(totalxyz, axis=0)
         rgb = np.concatenate(totalrgb, axis=0)
         totaltime = np.concatenate(totaltime, axis=0)
         assert xyz.shape[0] == rgb.shape[0]  
         xyzt =np.concatenate( (xyz, totaltime), axis=1)     
         storePly(totalply_path, xyzt, rgb)
     try:
+        print(f"Read merged point cloud file directly from: {totalply_path}")
         pcd = fetchPly(totalply_path)
     except:
         pcd = None

examples/scripts/dyn_exps/neu3dv/stg.sh

@@ -0,0 +1,48 @@
+SCENE_DIR="examples/data/neural_3d"
+
+SCENE_LIST="coffee_martini cook_spinach cut_roasted_beef flame_salmon_1 flame_steak sear_steak" # SCENE_LIST="coffee_martini cook_spinach cut_roasted_beef flame_salmon_1 flame_steak sear_steak"
+
+GPU_LIST=(0 1 2 3 4 5)
+
+RESULT_DIR="results/stg_neu3d"
+
+GOF_SIZE=50
+
+DEFAULT_GPU_ID=${GPU_LIST[0]}
+
+run_single_scene() {
+    local rd_lambda="$1"
+    local scene="$2"
+    local gpu_id="${3:-$DEFAULT_GPU_ID}"
+
+    # Training Phase
+    CUDA_VISIBLE_DEVICES=$gpu_id python examples/simple_trainer_STG.py default \
+        --model_path "$RESULT_DIR/default/$scene" \
+        --data_dir "$SCENE_DIR/$scene" \
+        --result_dir "$RESULT_DIR/default/$scene" \
+        --gof_num $GOF_SIZE \
+        --test_view_id 0 \
+        --disable_viewer \
+
+    # Compression Phase
+    CUDA_VISIBLE_DEVICES=$gpu_id python examples/simple_trainer_STG.py default \
+        --model_path "$RESULT_DIR/default/$scene" \
+        --data_dir "$SCENE_DIR/$scene" \
+        --result_dir "$RESULT_DIR/default/$scene" \
+        --gof_num $GOF_SIZE \
+        --test_view_id 0 \
+        --disable_viewer \
+        --compression_only \
+        --ckpt_name "ckpt_best_rank0.pt" \
+        --compression "stg"
+}
+
+
+idx=0
+for scene in $SCENE_LIST; do
+    gpu_id=${GPU_LIST[$((idx % ${#GPU_LIST[@]}))]}
+    run_single_scene "$rd_lambda" "$scene" "$gpu_id" &
+    idx=$((idx+1))
+done
+wait
+echo "Finished running all scenes."

examples/scripts/dyn_exps/neu3dv/stg_comp_sim.sh

@@ -0,0 +1,52 @@
+SCENE_DIR="examples/data/neural_3d"
+
+SCENE_LIST="coffee_martini cook_spinach cut_roasted_beef flame_salmon_1 flame_steak sear_steak" # SCENE_LIST="coffee_martini cook_spinach cut_roasted_beef flame_salmon_1 flame_steak sear_steak"
+
+GPU_LIST=(0 1 2 3 4 5)
+
+RESULT_DIR="results/stg_neu3d_comp_sim"
+
+GOF_SIZE=50
+
+RD_LAMBDA=(0.005 0.01 0.02) # refers to 3 ratepoints
+
+DEFAULT_GPU_ID=${GPU_LIST[0]}
+
+run_single_scene() {
+    local rd_lambda="$1"
+    local scene="$2"
+    local gpu_id="${3:-$DEFAULT_GPU_ID}"
+
+    # Training Phase
+    CUDA_VISIBLE_DEVICES=$gpu_id python examples/simple_trainer_STG.py compression_sim \
+        --model_path "$RESULT_DIR/rd_${rd_lambda}/$scene" \
+        --data_dir "$SCENE_DIR/$scene" \
+        --result_dir "$RESULT_DIR/rd_${rd_lambda}/$scene" \
+        --gof_num $GOF_SIZE \
+        --test_view_id 0 \
+        --disable_viewer \
+        --entropy_model_opt --rd_lambda $rd_lambda # enable entropy model with rd_lambda 0.01
+
+    # Compression Phase
+    CUDA_VISIBLE_DEVICES=$gpu_id python examples/simple_trainer_STG.py default \
+        --model_path "$RESULT_DIR/rd_${rd_lambda}/$scene" \
+        --data_dir "$SCENE_DIR/$scene" \
+        --result_dir "$RESULT_DIR/rd_${rd_lambda}/$scene" \
+        --gof_num $GOF_SIZE \
+        --test_view_id 0 \
+        --disable_viewer \
+        --compression_only \
+        --ckpt_name "ckpt_best_rank0.pt" \
+        --compression "stg"
+}
+
+for rd_lambda in "${RD_LAMBDA[@]}"; do
+    idx=0
+    for scene in $SCENE_LIST; do
+        gpu_id=${GPU_LIST[$((idx % ${#GPU_LIST[@]}))]}
+        run_single_scene "$rd_lambda" "$scene" "$gpu_id" &
+        idx=$((idx+1))
+    done
+    wait
+    echo "Finished running $rd_lambda on all scenes."
+done