support replica

Author: yuepan

.gitignore

@@ -5,4 +5,6 @@ __pycache__/
 *.txt
 *.ply
 *.bin
-*.pth
+*.pth
+./experiments/
+./data/

README.md

@@ -245,6 +245,10 @@ For the sake of efficiency, we sacrifice a bit mapping quality to use a 50cm lea
   <img src="https://user-images.githubusercontent.com/34207278/216335116-45273ef9-adb8-4f03-9e58-c5ca60b03081.png" width="70%" />
 </p>
 
+<p align="center">
+  <img src="https://user-images.githubusercontent.com/34207278/217082027-efe272bb-ff6e-49b3-bb69-6d0a32c64a45.png" width="70%" />
+</p>
+
 **Newer College**
 
 <p align="center">
@@ -257,24 +261,32 @@ For the sake of efficiency, we sacrifice a bit mapping quality to use a 50cm lea
   <img src="https://user-images.githubusercontent.com/34207278/212333973-ad1a0c14-550f-4266-902f-08c33f9f79c1.png" width="70%" />
 </p>
 
-**RGB-D**
-
-<p align="center">
-  <img src="https://user-images.githubusercontent.com/34207278/206718688-a8bfb786-6439-435b-a225-c596aef9ab64.png" width="70%" />
-</p>
 
 **Wild Place Forests**
 
 <p align="center">
   <img src="https://user-images.githubusercontent.com/34207278/207911217-c27a52c9-7233-4db9-a1fd-3487e59e6529.png" width="70%" />
 </p>
 
-**Office**
+
+**IPB Office**
 
 <p align="center">
   <img src="https://user-images.githubusercontent.com/34207278/216360654-9b0a8bda-6a98-4db1-aa25-1c58080a4585.png" width="70%" />
 </p>
 
+**ICL Living Room**
+
+<p align="center">
+  <img src="https://user-images.githubusercontent.com/34207278/206718688-a8bfb786-6439-435b-a225-c596aef9ab64.png" width="70%" />
+</p>
+
+**Replica**
+
+<p align="center">
+  <img src="https://user-images.githubusercontent.com/34207278/217082679-319d220a-de1f-4f30-ac14-438b4fe8942a.png" width="70%" />
+</p>
+
 
 </details>
 

config/kitti/kitti_batch.yaml

@@ -53,10 +53,11 @@ optimizer:
   weight_decay: 1e-7 # l2 regularization
 eval:
   wandb_vis_on: False # log to wandb or not
+  o3d_vis_on: True # visualize the mapping or not
   vis_freq_iters: 5000 # reconstruct the mesh every x iterations
   save_freq_iters: 10000 # save the model and octree every x iterations
   mesh_freq_frame: 1
   mc_res_m: 0.1 # reconstruction marching cubes resolution
-  mc_mask_on: True
+  mc_with_octree: True # querying sdf only in certain levels of the octree
   mc_vis_level: 1
   save_map: False # save the sdf map or not 

config/kitti/kitti_incre_reg.yaml

@@ -53,10 +53,11 @@ optimizer:
   weight_decay: 1e-7 # l2 regularization
 eval:
   wandb_vis_on: False # log to wandb or not
+  o3d_vis_on: True # visualize the mapping or not
   vis_freq_iters: 0 
   save_freq_iters: 10000 # save the model and octree every x iterations
   mesh_freq_frame: 5 # reconstruct the mesh every x frames
   mc_res_m: 0.1 # reconstruction marching cubes resolution
-  mc_mask_on: True
+  mc_with_octree: True # querying sdf only in certain levels of the octree
   mc_vis_level: 2
   save_map: False # save the sdf map or not 

config/kitti/kitti_incre_replay.yaml

@@ -53,10 +53,11 @@ optimizer:
   weight_decay: 1e-7 # l2 regularization
 eval:
   wandb_vis_on: False # log to wandb or not
+  o3d_vis_on: True # visualize the mapping or not
   vis_freq_iters: 0 
   save_freq_iters: 10000 # save the model and octree every x iterations
   mesh_freq_frame: 5 # reconstruct the mesh every x frames
   mc_res_m: 0.1 # reconstruction marching cubes resolution
-  mc_mask_on: True
+  mc_with_octree: True # querying sdf only in certain levels of the octree
   mc_vis_level: 2
   save_map: False # save the sdf map or not 

config/maicity/maicity_batch.yaml

@@ -53,10 +53,11 @@ optimizer:
   weight_decay: 1e-7 # l2 regularization
 eval:
   wandb_vis_on: False # log to wandb or not
+  o3d_vis_on: True # visualize the mapping or not
   vis_freq_iters: 5000 # reconstruct the mesh every x iterations
   save_freq_iters: 10000 # save the model and octree every x iterations
   mesh_freq_frame: 1 
   mc_res_m: 0.1 # reconstruction marching cubes resolution
-  mc_mask_on: True
+  mc_with_octree: False # querying sdf in the map bbx
   mc_vis_level: 1 # mask the voxels used for marching cubes reconstruction according to the octree level
   save_map: False # save the sdf map or not

config/maicity/maicity_incre_reg.yaml

@@ -53,10 +53,11 @@ optimizer:
   weight_decay: 0 # l2 regularization
 eval:
   wandb_vis_on: False # log to wandb or not
+  o3d_vis_on: True # visualize the mapping or not
   vis_freq_iters: 0 
   save_freq_iters: 0 # save the model and octree every x iterations
   mesh_freq_frame: 5 # reconstruct the mesh every x frames
   mc_res_m: 0.2 # reconstruction marching cubes resolution
-  mc_mask_on: True
+  mc_with_octree: False # querying sdf in the map bbx
   mc_vis_level: 2
   save_map: False # save the sdf map or not

config/maicity/maicity_incre_replay.yaml

@@ -53,10 +53,11 @@ optimizer:
   weight_decay: 0 # l2 regularization
 eval:
   wandb_vis_on: False # log to wandb or not
+  o3d_vis_on: True # visualize the mapping or not
   vis_freq_iters: 0 
   save_freq_iters: 0 # save the model and octree every x iterations
   mesh_freq_frame: 5 # reconstruct the mesh every x frames
   mc_res_m: 0.2 # reconstruction marching cubes resolution
-  mc_mask_on: True
+  mc_with_octree: False # querying sdf in the map bbx
   mc_vis_level: 2 
   save_map: False # save the sdf map or not

config/ncd/ncd_batch.yaml

@@ -53,10 +53,11 @@ optimizer:
   weight_decay: 1e-7 # l2 regularization
 eval:
   wandb_vis_on: False # log to wandb or not
+  o3d_vis_on: True # visualize the mapping or not
   vis_freq_iters: 5000 # reconstruct the mesh every x iterations
   save_freq_iters: 10000 # save the model and octree every x iterations
   mesh_freq_frame: 1
   mc_res_m: 0.1 # reconstruction marching cubes resolution
-  mc_mask_on: True
+  mc_with_octree: True # querying sdf only in certain levels of the octree
   mc_vis_level: 1
   save_map: False # save the sdf map or not 

config/ncd/ncd_incre_reg.yaml

@@ -53,10 +53,11 @@ optimizer:
   weight_decay: 1e-7 # l2 regularization
 eval:
   wandb_vis_on: False # log to wandb or not
+  o3d_vis_on: True # visualize the mapping or not
   vis_freq_iters: 0 
   save_freq_iters: 10000 # save the model and octree every x iterations
   mesh_freq_frame: 5 # reconstruct the mesh every x frames
   mc_res_m: 0.1 # reconstruction marching cubes resolution
-  mc_mask_on: True
+  mc_with_octree: True # querying sdf only in certain levels of the octree
   mc_vis_level: 2
   save_map: False # save the sdf map or not 

config/ncd/ncd_incre_replay.yaml

@@ -53,10 +53,11 @@ optimizer:
   weight_decay: 1e-7 # l2 regularization
 eval:
   wandb_vis_on: False # log to wandb or not
+  o3d_vis_on: True # visualize the mapping or not
   vis_freq_iters: 0 
   save_freq_iters: 10000 # save the model and octree every x iterations
   mesh_freq_frame: 5 # reconstruct the mesh every x frames
   mc_res_m: 0.1 # reconstruction marching cubes resolution
-  mc_mask_on: True
+  mc_with_octree: True # querying sdf only in certain levels of the octree
   mc_vis_level: 2
   save_map: False # save the sdf map or not 

config/rgbd/rgbd_batch.yaml

@@ -53,10 +53,11 @@ optimizer:
   weight_decay: 1e-7 # l2 regularization
 eval:
   wandb_vis_on: False # log to wandb or not
+  o3d_vis_on: True # visualize the mapping or not
   vis_freq_iters: 5000 # reconstruct the mesh every x iterations
   save_freq_iters: 10000 # save the model and octree every x iterations
   mesh_freq_frame: 1
-  mc_res_m: 0.02 # reconstruction marching cubes resolution
-  mc_mask_on: True
+  mc_res_m: 0.01 # reconstruction marching cubes resolution
+  mc_with_octree: True # querying sdf in the map bbx
   mc_vis_level: 2
   save_map: False # save the sdf map or not 

config/rgbd/rgbd_incre_reg.yaml

@@ -53,10 +53,11 @@ optimizer:
   weight_decay: 1e-7 # l2 regularization
 eval:
   wandb_vis_on: False # log to wandb or not
+  o3d_vis_on: True # visualize the mapping or not
   vis_freq_iters: 0 # reconstruct the mesh every x iterations
   save_freq_iters: 10000 # save the model and octree every x iterations
   mesh_freq_frame: 5 # reconstruct the mesh every x frames
   mc_res_m: 0.02 # reconstruction marching cubes resolution
-  mc_mask_on: True
+  mc_with_octree: True # querying sdf in the map bbx
   mc_vis_level: 2
   save_map: False # save the sdf map or not 

config/rgbd/rgbd_incre_replay.yaml

@@ -53,10 +53,11 @@ optimizer:
   weight_decay: 1e-7 # l2 regularization
 eval:
   wandb_vis_on: False # log to wandb or not
+  o3d_vis_on: True # visualize the mapping or not
   vis_freq_iters: 0 # reconstruct the mesh every x iterations
   save_freq_iters: 10000 # save the model and octree every x iterations
   mesh_freq_frame: 5 # reconstruct the mesh every x frames
   mc_res_m: 0.02 # reconstruction marching cubes resolution
-  mc_mask_on: True
+  mc_with_octree: True # querying sdf in the map bbx
   mc_vis_level: 2
   save_map: False # save the sdf map or not 

dataset/rgbd_to_kitti_format.py

@@ -2,21 +2,23 @@
 from tqdm import tqdm
 import argparse
 import re
-import json
 import os
 import numpy as np
-
-from utils.tools import load_from_json
+import json
+import shutil
 
 def rgbd_to_kitti_format(args):
 
     ply_path = os.path.join(args.output_root, "rgbd_ply")
     os.makedirs(ply_path, exist_ok=True)
 
     # get pose
-    poses_mat = load_poses(args.pose_file, with_head = False) # with_head = True for open3d provided Redwood dataset
     pose_kitti_format_path = os.path.join(args.output_root, "poses.txt")
-    write_poses_kitti_format(poses_mat, pose_kitti_format_path)
+    if args.already_kitti_format_pose:
+        shutil.copyfile(args.pose_file, pose_kitti_format_path) # don't directly copy, may have some issues
+    else:
+        poses_mat = load_poses(args.pose_file, with_head = False) # with_head = True for open3d provided Redwood dataset
+        write_poses_kitti_format(poses_mat, pose_kitti_format_path)
 
     # get an example image 
     depth_img_files = sorted(os.listdir(args.depth_img_folder), key=alphanum_key)
@@ -36,8 +38,10 @@ def rgbd_to_kitti_format(args):
         # W=640, H=480, fx=fy=525.0, cx=319.5, cy=239.5
         print("Default intrinsic for PrimeSense used")
         intrinsic = o3d.camera.PinholeCameraIntrinsic(o3d.camera.PinholeCameraIntrinsicParameters.PrimeSenseDefault)
+        # use this extrinsic matrix to rotate the image since frames captured with RealSense camera are upside down 
+        extrinsic = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
     else:
-        if args.is_focal_file: # load the focal length only txt file
+        if args.is_focal_file: # load the focal length only txt file # This is used for NeuralRGBD dataset
             focal = load_focal_length(args.intrinsic_file)
             print("Focal length:", focal)
             intrinsic.set_intrinsics(height=H,
@@ -47,6 +51,8 @@ def rgbd_to_kitti_format(args):
                                      cx=(W-1.)/2.,
                                      cy=(H-1.)/2.)
             depth_scale = 1000.
+            # use this extrinsic matrix to rotate the image since frames captured with RealSense camera are upside down 
+            extrinsic = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
 
         else:
             with open(args.intrinsic_file, 'r') as infile: # load intrinsic json file
@@ -58,9 +64,8 @@ def rgbd_to_kitti_format(args):
                                         cx=cam["cx"],
                                         cy=cam["cy"])
                 depth_scale = cam["scale"]
+                extrinsic = np.eye(4) # this is used for Replica dataset
 
-    # use this extrinsic matrix to rotate the image since frames captured with RealSense camera are upside down 
-    extrinsic = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
 
     # get point cloud
     frame_count = 0
@@ -69,9 +74,10 @@ def rgbd_to_kitti_format(args):
         depth_path = os.path.join(args.depth_img_folder, depth_path)
         print(color_path)
         print(depth_path)
+        
         im_color = o3d.io.read_image(color_path)
-        im_depth = o3d.io.read_image(depth_path)
-        im_rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(im_color, im_depth, depth_scale = depth_scale, depth_trunc = args.max_depth_m)
+        im_depth = o3d.io.read_image(depth_path) 
+        im_rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(im_color, im_depth, depth_scale = depth_scale, depth_trunc = args.max_depth_m, convert_rgb_to_intensity=False) # not just gray
         im_pcd = o3d.geometry.PointCloud.create_from_rgbd_image(im_rgbd, intrinsic, extrinsic)
 
         if args.vis_on:
@@ -101,6 +107,15 @@ def alphanum_key(s):
     """
     return [int(x) if x.isdigit() else x for x in re.split('([0-9]+)', s)]
 
+def load_from_json(filename):
+    """Load a dictionary from a JSON filename.
+    Args:
+        filename: The filename to load from.
+    """
+    assert filename.suffix == ".json"
+    with open(filename, encoding="UTF-8") as file:
+        return json.load(file)
+
 def load_focal_length(filepath):
     file = open(filepath, "r")
     return float(file.readline())
@@ -147,6 +162,9 @@ def parser_json_sdf_studio_format(json_file):
     parser.add_argument('--is_focal_file', type=str2bool, nargs='?', default=True, \
         help="is the input intrinsic file a txt file containing only the focus length (as the Neural RGBD data format)\
               or the json file containing all the intrinsic parameters (as the Replica format)")
+    parser.add_argument('--already_kitti_format_pose', type=str2bool, nargs='?', default=False, \
+        help="is the input pose file already in KITTI pose format (also as the Replica format)\
+              or the input pose file is in a 4dim transformation form (as the Neural RGBD data format)")
     parser.add_argument('--vis_on', type=str2bool, nargs='?', default=False)
     args = parser.parse_args()
 

model/feature_octree.py

@@ -1,19 +1,15 @@
 import torch
-import torch.nn.functional as F
 import torch.nn as nn
 
-import math
 import time
 from tqdm import tqdm
 import kaolin as kal
 import numpy as np
-from torch.autograd import grad
 
 from functools import partial
 from collections import defaultdict
 
 from utils.config import SHINEConfig
-from utils.loss import sdf_bce_loss
 
 # TODO: polish the codes
 

scripts/convert_rgbd_to_kitti_format.sh

@@ -4,13 +4,17 @@ root_path=xxx/data/neural_rgbd_data
 sequence_name=green_room
 base_path=${root_path}/${sequence_name}
 
+# For NeuralRGBD dataset, set is_focal_file to True, and already_kitti_format_pose to False
+# For Replica dataset,    set is_focal_file to False, and already_kitti_format_pose to True
+
 command="python3 ./dataset/rgbd_to_kitti_format.py
         --output_root ${base_path}_kitti_format
         --depth_img_folder ${base_path}/depth_filtered/
         --rgb_img_folder ${base_path}/images/
         --intrinsic_file ${base_path}/focal.txt
         --pose_file ${base_path}/poses.txt
         --is_focal_file True
+        --already_kitti_format_pose False
         --vis_on False"
 
 echo "Convert RGBD dataset to KITTI format"

shine_incre.py

@@ -162,6 +162,7 @@ def run_shine_mapping_incremental():
 
         # reconstruction by marching cubes
         if processed_frame == 0 or (processed_frame+1) % config.mesh_freq_frame == 0:
+            print("Begin mesh reconstruction from the implicit map")       
             vis_mesh = True 
             # print("Begin reconstruction from implicit mapn")               
             mesh_path = run_path + '/mesh/mesh_frame_' + str(frame_id+1) + ".ply"