diff --git a/.gitmodules b/.gitmodules index 90586ea..3df6822 100644 --- a/.gitmodules +++ b/.gitmodules @@ -10,5 +10,5 @@ branch = feature/nerf_slam [submodule "thirdparty/gtsam"] path = thirdparty/gtsam - url = https://github.com/ToniRV/gtsam-1.git - branch = feature/nerf_slam + url = https://github.com/borglab/gtsam.git + branch = develop diff --git a/README.md b/README.md index a3337a1..bf8a4d6 100644 --- a/README.md +++ b/README.md @@ -70,28 +70,39 @@ Clone repo with submodules: ``` git clone https://github.com/ToniRV/NeRF-SLAM.git --recurse-submodules git submodule update --init --recursive +cd thirdparty/instant-ngp/ && git checkout feature/nerf_slam ``` From this point on, use a virtual environment... Install torch (see [here](https://pytorch.org/get-started/previous-versions) for other versions): + +### Install CUDA 11.7 and PyTorch + +Manually install [CUDA 11.7 here](https://developer.nvidia.com/cuda-11-7-1-download-archive). + +Or, if using conda: +``` +conda install -c "nvidia/label/cuda-11.7.0" cuda-toolkit +``` +Then install pytorch: ``` -# CUDA 11.3 -pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113 +pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117 ``` -Pip install requirements: + +### Pip install requirements: ``` pip install -r requirements.txt pip install -r ./thirdparty/gtsam/python/requirements.txt ``` -Compile ngp (you need cmake>3.22): +### Compile ngp (you need cmake>3.22): ``` cmake ./thirdparty/instant-ngp -B build_ngp cmake --build build_ngp --config RelWithDebInfo -j ``` -Compile gtsam and enable the python wrapper: +### Compile gtsam and enable the python wrapper: ``` cmake ./thirdparty/gtsam -DGTSAM_BUILD_PYTHON=1 -B build_gtsam cmake --build build_gtsam --config RelWithDebInfo -j @@ -99,7 +110,7 @@ cd build_gtsam make python-install ``` -Install: +### Install: ``` python setup.py install ``` @@ -119,6 +130,14 @@ python ./examples/slam_demo.py --dataset_dir=./datasets/Replica/office0 --datase This repo also implements [Sigma-Fusion](https://arxiv.org/abs/2210.01276): just change `--fusion='sigma'` to run that. +### Other Run modes + +Skip SLAM, use GT poses and depth with the cube diorama scene: +``` +./scripts/download_cube.bash +python ./examples/slam_demo.py --dataset_dir=./datasets/nerf-cube-diorama-dataset/room --dataset_name=nerf --buffer=100 --img_stride=1 --fusion='nerf' --gui +``` + ## FAQ ### GPU Memory diff --git a/datasets/replica_dataset.py b/datasets/replica_dataset.py index 34b4b74..39b7c53 100644 --- a/datasets/replica_dataset.py +++ b/datasets/replica_dataset.py @@ -5,7 +5,7 @@ import numpy as np import cv2 -import tqdm +from tqdm import tqdm from icecream import ic from datasets.dataset import * @@ -66,18 +66,52 @@ def parse_dataset(self): N = self.args.buffer H, W = self.calib.resolution.height, self.calib.resolution.width + self.resize_images = False + if self.calib.resolution.total() > 640*640: + self.resize_images = True + # TODO(Toni): keep aspect ratio, and resize max res to 640 + self.output_image_size = [341, 640] # h, w + + if self.resize_images: + h0, w0 = self.calib.resolution.height, self.calib.resolution.width + total_output_pixels = (self.output_image_size[0] * self.output_image_size[1]) + self.h1 = int(h0 * np.sqrt(total_output_pixels / (h0 * w0))) + self.w1 = int(w0 * np.sqrt(total_output_pixels / (h0 * w0))) + self.h1 = self.h1 - self.h1 % 8 + self.w1 = self.w1 - self.w1 % 8 + self.calib.camera_model.scale_intrinsics(self.w1 / w0, self.h1 / h0) + self.calib.resolution = Resolution(self.w1, self.h1) + + subset_poses = [] + + if self.final_k == -1: + self.final_k = len(self.poses) - 1 + # Parse images and tfs for i, (image_path, depth_path) in enumerate(tqdm(zip(self.image_paths, self.depth_paths))): - if i >= N: - break + + if ((i-self.initial_k) % self.img_stride) != 0 or i < self.initial_k or i > self.final_k: + continue # Parse rgb/depth images image = cv2.imread(image_path) depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED) - image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # this is for NERF - depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)[..., None] # H, W, C=1 + #depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)[..., None] # H, W, C=1 + if self.resize_images: + w1, h1 = self.w1, self.h1 + image = cv2.resize(image, (w1, h1)) + image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA) # Required for Nerf Fusion, perhaps we can put it in there + + depth = cv2.resize(depth, (w1, h1)) + depth = depth[:, :, np.newaxis] + + if self.viz: + cv2.imshow(f"Img Resized", image) + cv2.imshow(f"Depth Resized", depth) + cv2.waitKey(1) + H, W, _ = depth.shape assert(H == image.shape[0]) assert(W == image.shape[1]) @@ -94,8 +128,13 @@ def parse_dataset(self): self.images += [image] self.depths += [depth] self.calibs += [self.calib] + subset_poses += [self.poses[i]] + + # Early break if we've exceeded the buffer max + if len(self.images) == self.args.buffer: + break - self.poses = self.poses[:N] + self.poses = subset_poses self.timestamps = np.array(self.timestamps) self.poses = np.array(self.poses) diff --git a/examples/slam_demo.py b/examples/slam_demo.py index 45245de..3b3f5ac 100644 --- a/examples/slam_demo.py +++ b/examples/slam_demo.py @@ -110,12 +110,14 @@ def run(args): data_provider_module.register_output_queue(data_for_fusion_output_queue) fusion_module.register_input_queue("data", data_for_fusion_output_queue) + # Create interactive Gui gui = args.gui and args.fusion != 'nerf' # nerf has its own gui if gui: gui_module = GuiModule("Open3DGui", args, device=cuda_slam) # don't use cuda:1, o3d doesn't work... data_provider_module.register_output_queue(data_for_viz_output_queue) - slam_module.register_output_queue(slam_output_queue_for_o3d) + if slam: + slam_module.register_output_queue(slam_output_queue_for_o3d) gui_module.register_input_queue("data", data_for_viz_output_queue) gui_module.register_input_queue("slam", slam_output_queue_for_o3d) if fusion and (fusion_module.name == "tsdf" or fusion_module.name == "sigma"): diff --git a/fusion/nerf_fusion.py b/fusion/nerf_fusion.py index 4106c35..33d8ffe 100644 --- a/fusion/nerf_fusion.py +++ b/fusion/nerf_fusion.py @@ -214,7 +214,7 @@ def process_slam(self, packet): else: images = srgb_to_linear(images, self.device) - data_packets = {"k": viz_idx, + data_packets = {"k": viz_idx.cpu().numpy(), "poses": world_T_cam0, # needs to be c2w "images": images.contiguous().cpu().numpy(), "depths": depths.contiguous().cpu().numpy(), @@ -282,12 +282,17 @@ def send_data(self, batch): principal_point = intrinsics[2:] # TODO: we need to restore the self.ref_frames[frame_id] = [image, gt, etc] for evaluation.... - self.ngp.nerf.training.update_training_images(frame_ids.cpu().numpy().tolist(), + self.ngp.nerf.training.update_training_images(list(frame_ids), list(poses[:, :3, :4]), list(images), list(depths), list(depths_cov), resolution, principal_point, focal_length, depth_scale, depth_cov_scale) + # On the first frame, set the viewpoint + if self.ngp.nerf.training.n_images_for_training == 1: + self.ngp.set_camera_to_training_view(0) + + def fit_volume(self): #print(f"Fitting volume for {self.iters} iters") self.fps = 30 diff --git a/requirements.txt b/requirements.txt index d10dc2d..59b1b0b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -19,3 +19,5 @@ pyrealsense2 pybind11 gdown + +colored_glog \ No newline at end of file diff --git a/scripts/download_cube.bash b/scripts/download_cube.bash index 0efe968..8c8592f 100755 --- a/scripts/download_cube.bash +++ b/scripts/download_cube.bash @@ -1,5 +1,5 @@ #!/usr/bin/env sh -mkdir -p Datasets -cd Datasets +mkdir -p datasets +cd datasets git clone https://github.com/jc211/nerf-cube-diorama-dataset.git diff --git a/scripts/download_replica.bash b/scripts/download_replica.bash index 2ac254f..74c0faf 100755 --- a/scripts/download_replica.bash +++ b/scripts/download_replica.bash @@ -1,7 +1,7 @@ #!/usr/bin/env sh -mkdir -p Datasets -cd Datasets +mkdir -p datasets +cd datasets # This is 9.2Gb of data: contains all the images, transforms (.json), and meshes (.ply). gdown 1tdioYdNGK6yZZdfyQKkQI84ALtEqd2fH diff --git a/scripts/download_replica_sample.bash b/scripts/download_replica_sample.bash index 895f526..733a039 100755 --- a/scripts/download_replica_sample.bash +++ b/scripts/download_replica_sample.bash @@ -1,10 +1,10 @@ #!/usr/bin/env sh -mkdir -p Datasets -cd Datasets +mkdir -p datasets +cd datasets gdown 1f4RJ4W9uxlhCvihG12X9Wa4yZrlBD4TE mkdir -p Replica/ -unzip ReplicaSample.zip -d Replica +unzip replica_sample.zip -d Replica diff --git a/slam/inertial_frontends/inertial_frontend.py b/slam/inertial_frontends/inertial_frontend.py index 6bc591a..a264652 100644 --- a/slam/inertial_frontends/inertial_frontend.py +++ b/slam/inertial_frontends/inertial_frontend.py @@ -172,7 +172,7 @@ def initial_state(self): gtsam.Point3(0.878612,2.142470,0.947262)) true_vel = np.array([0.009474,-0.014009,-0.002145]) true_bias = gtsam.imuBias.ConstantBias(np.array([-0.012492,0.547666,0.069073]), np.array([-0.002229,0.020700,0.076350])) - naive_pose = gtsam.Pose3.identity() + naive_pose = gtsam.Pose3() #identity naive_vel = np.zeros(3) naive_bias = gtsam.imuBias.ConstantBias() initial_pose = true_pose diff --git a/slam/vio_slam.py b/slam/vio_slam.py index 5e0f3ac..85a6550 100644 --- a/slam/vio_slam.py +++ b/slam/vio_slam.py @@ -62,7 +62,7 @@ def initial_state(): gtsam.Point3(0.878612, 2.142470, 0.947262)) true_vel = np.array([0.009474,-0.014009,-0.002145]) true_bias = gtsam.imuBias.ConstantBias(np.array([-0.012492,0.547666,0.069073]), np.array([-0.002229,0.020700,0.076350])) - naive_pose = gtsam.Pose3.identity() + naive_pose = gtsam.Pose3() #identity naive_vel = np.zeros(3) naive_bias = gtsam.imuBias.ConstantBias() initial_pose = true_world_T_imu_t0 diff --git a/slam/visual_frontends/visual_frontend.py b/slam/visual_frontends/visual_frontend.py index 5b1a75b..4c630af 100644 --- a/slam/visual_frontends/visual_frontend.py +++ b/slam/visual_frontends/visual_frontend.py @@ -45,8 +45,8 @@ def lietorch_pose_to_gtsam(pose : lietorch.SE3): def gtsam_pose_to_torch(pose: gtsam.Pose3, device, dtype): t = pose.translation() - q = pose.rotation().quaternion() - return torch.tensor([t[0], t[1], t[2], q[1], q[2], q[3], q[0]], device=device, dtype=dtype) + q = pose.rotation().toQuaternion() + return torch.tensor([t[0], t[1], t[2], q.x(), q.y(), q.z(), q.w()], device=device, dtype=dtype) class VisualFrontend(nn.Module): def __init__(self): diff --git a/thirdparty/gtsam b/thirdparty/gtsam index 139c156..f427b04 160000 --- a/thirdparty/gtsam +++ b/thirdparty/gtsam @@ -1 +1 @@ -Subproject commit 139c156a7a6fb5ef6062549993410a8ab986720b +Subproject commit f427b0411b0118c9e28edc9fbd73c3e01efd911d diff --git a/thirdparty/instant-ngp b/thirdparty/instant-ngp index 54aba7c..a501f0f 160000 --- a/thirdparty/instant-ngp +++ b/thirdparty/instant-ngp @@ -1 +1 @@ -Subproject commit 54aba7cfbeaf6a60f29469a9938485bebeba24c3 +Subproject commit a501f0f14699a2b493f1e2fd42a975a8a9c172d1