diff --git a/.gitmodules b/.gitmodules index 63b426a6..1d3fdfd3 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,7 @@ [submodule "diff-gaussian-rasterization-w-depth.git"] path = diff-gaussian-rasterization-w-depth.git url = git@github.com:JonathonLuiten/diff-gaussian-rasterization-w-depth.git +[submodule "spectacularAI-sdk"] + path = spectacularAI-sdk + url = https://github.com/SpectacularAI/sdk-examples + branch = nerfcapture-export diff --git a/README.md b/README.md index 766d6a9b..261018d2 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,7 @@ Installation
  • - Online Demo + Offline Demo
  • Usage @@ -66,7 +66,6 @@ ## Installation -##### (Recommended) SplaTAM has been benchmarked with Python 3.10, Torch 1.12.1 & CUDA=11.6. However, Torch 1.12 is not a hard requirement and the code has also been tested with other versions of Torch and CUDA such as Torch 2.3.0 & CUDA 12.1. The simplest way to install all dependences is to use [anaconda](https://www.anaconda.com/) and [pip](https://pypi.org/project/pip/) in the following steps: @@ -86,6 +85,7 @@ conda activate splatam ``` --> #### Docker Setup +##### (Ekumen Recommended) We also provide a docker image. We recommend using a venv to run the code inside a docker image: @@ -101,60 +101,38 @@ pip install -r venv_requirements.txt ## Demo -### Online - -You can SplaTAM your own environment with an iPhone or LiDAR-equipped Apple device by downloading and using the NeRFCapture app. - -Make sure that your iPhone and PC are connected to the same WiFi network, and then run the following command: - - ```bash -bash bash_scripts/online_demo.bash configs/iphone/online_demo.py -``` - -On the app, keep clicking send for successive frames. Once the capturing of frames is done, the app will disconnect from the PC and check out SplaTAM's interactive rendering of the reconstruction on your PC! Here are some cool example results: - -

    - - Logo - -

    - ### Offline -You can also first capture the dataset and then run SplaTAM offline on the dataset with the following command: +You can first capture the dataset with Spectacular AI tool and then run SplaTAM offline on the recorded dataset with the following command: ```bash -bash bash_scripts/nerfcapture.bash configs/iphone/nerfcapture.py +bash bash_scripts/nerfcapture.bash configs/iphone/nerfcapture.py ``` ### Dataset Collection -If you would like to only capture your own iPhone dataset using the NeRFCapture app, please use the following command: - -```bash -bash bash_scripts/nerfcapture2dataset.bash configs/iphone/dataset.py -``` +If you would like to only capture your own iPhone dataset, please use the Spectacular AI app. ## Usage -We will use the iPhone dataset as an example to show how to use SplaTAM. The following steps are similar for other datasets. - -To run SplaTAM, please use the following command: +If you already converted the dataset to the SplaTAM expected input, to run SplaTAM, please use the following command: ```bash -python scripts/splatam.py configs/iphone/splatam.py +python scripts/splatam.py configs/iphone/splatam.py --base_dir --scene --frames ``` To visualize the final interactive SplaTAM reconstruction, please use the following command: ```bash -python viz_scripts/final_recon.py configs/iphone/splatam.py +python viz_scripts/final_recon.py configs/iphone/splatam.py --base_dir --scene --frames ``` +### Not supported by Ekumen + To visualize the SplaTAM reconstruction in an online fashion, please use the following command: ```bash -python viz_scripts/online_recon.py configs/iphone/splatam.py +python viz_scripts/online_recon.py configs/iphone/splatam.py --base_dir --scene --frames ``` To export the splats to a .ply file, please use the following command: diff --git a/RESULTS.md b/RESULTS.md new file mode 100644 index 00000000..5cad64b7 --- /dev/null +++ b/RESULTS.md @@ -0,0 +1,71 @@ +# 3D Mapping Workflow (Frodo) + +# Recording devices + +List of the recording devices used during the Frodo + +* iPhone Pro + * Has lidar, which allows depth dataset recording. +* iPhone + * No lidar, no depth capture. +* Android device + * ARCore should suffice for the missing lidar, but there was no depth dataset saved after recording. + * So no depth capture. + +# Recording tools + +## NeRF Capture + +* Only available for IOS. +* Provides 2 different recording modes: + * Online mode & offline mode. + +### Online mode + +Requires the host to be connected into the recording device, so as to process the input while it is being captured. + +### Offline mode + +Saves a directory with the RBG captured frames and the depth frames \+ a transform.json file to feed into the processing tool. +Offline mode is currently bugged. There is an issue with the recording processing, which is not converting the depth buffer correctly into 16 bit numbers. See [this author comment](https://github.com/jc211/NeRFCapture/issues/10#issuecomment-1888164311) for more details. + +## Spectacular AI + +* Available for both Android and IOS. +* IOS’ version provides significantly less configuration options than Android’s. + * But Android devices don’t generate depth dataset. + +Recording should be taken by moving through a room while pointing the device towards the walls, and occasionally performing a panoramic view of the room to capture objects that are centered or away from the walls. The open3d capture from the ping pong table shows a good example of a trajectory the recorder should follow. + +# Processing tools + +## SplaTAM + +### Spectacular AI dataset + +In order to run SplaTAM with Spectacular AI recorded data, the first step was to process it to match the output format of NeRF Capture. + +### NeRF Capture dataset + +To get the results, SplaTAM provides the [nerfcapture2dataset](https://github.com/ekumenlabs/SplaTAM/blob/main/scripts/nerfcapture2dataset.py) script to process the NeRF Capture output into the expected directory structure. +This script was modified to work with the Offline mode, by allowing to load a previously recorded dataset + +After running [splatam](https://github.com/ekumenlabs/SplaTAM/blob/main/scripts/splatam.py) algorithm, the output is saved into an params.npz file format that can then be visualized with open3d viewer to get ahold of the desired point cloud, alongside a metrics plot showing the PSNR and the captured depth of each frame, and additionally, each output frame rasterized. +See [NpzFile](https://numpy.org/devdocs/reference/generated/numpy.lib.npyio.NpzFile.html#numpy.lib.npyio.NpzFile) for info on `npz` format and how to load it. + +### Results + +First Spectacular AI results show metrics plots that look good on the first third of the processed frames, but derailed on the rest. This behavior was repeated throughout 3 captured datasets with an iPhone 14 pro. +![image](assets/metrics.png) +This leads to the rasterized images and the point cloud visualization to look awful, since the tracking of the capturing device’s position is not correct. + +![image](assets/hamilton-broken.png) + +Changing the amount of frames to be processed to the point where the PSNR drops shows a much better result. It can be seen how the device tracker is within the captured environment. +![image](assets/hamilton-02-open3d.png) +![image](assets/pool-table-open3d.png) + +## Instant NGP + +\[TODO\] + diff --git a/assets/hamilton-02-open3d.png b/assets/hamilton-02-open3d.png new file mode 100644 index 00000000..6ac2e2b3 Binary files /dev/null and b/assets/hamilton-02-open3d.png differ diff --git a/assets/hamilton-broken.png b/assets/hamilton-broken.png new file mode 100644 index 00000000..4b5cefb7 Binary files /dev/null and b/assets/hamilton-broken.png differ diff --git a/assets/metrics.png b/assets/metrics.png new file mode 100644 index 00000000..f510236e Binary files /dev/null and b/assets/metrics.png differ diff --git a/assets/pool-table-open3d.png b/assets/pool-table-open3d.png new file mode 100644 index 00000000..d0ca8f4f Binary files /dev/null and b/assets/pool-table-open3d.png differ diff --git a/bash_scripts/nerfcapture.bash b/bash_scripts/nerfcapture.bash index 9dd8b118..d07f3374 100644 --- a/bash_scripts/nerfcapture.bash +++ b/bash_scripts/nerfcapture.bash @@ -1,8 +1,8 @@ #!/bin/bash # check rmem_max and wmem_max, and increase size if necessary -if [ "$#" -ne 1 ]; then - echo "Usage: bash_scripts/nerfcapture.bash " +if [ "$#" -ne 2 ]; then + echo "Usage: bash_scripts/nerfcapture.bash " exit fi @@ -15,21 +15,43 @@ if sysctl -a | grep -q "net.core.rmem_max = 2147483647"; then echo "rmem_max already set to 2147483647" else echo "Setting rmem_max to 2147483647" - sudo sysctl -w net.core.rmem_max=2147483647 + sysctl -w net.core.rmem_max=2147483647 fi if sysctl -a | grep -q "net.core.wmem_max = 2147483647"; then echo "wmem_max already set to 2147483647" else echo "Setting wmem_max to 2147483647" - sudo sysctl -w net.core.wmem_max=2147483647 + sysctl -w net.core.wmem_max=2147483647 fi -# Capture Dataset -python3 scripts/nerfcapture2dataset.py --config $1 +base_dir=$(dirname $2) +scene="$(basename $2)_nerfcapture" +echo $base_dir $scene + +# Convert Spectacular AI dataset to NeRF Capture. +echo "Exporting dataset to NeRF Capture format..." +python3 spectacularAI-sdk/python/mapping/replay_to_nerf.py $2 \ + --format=nerfcapture --fast \ + --image_format=png \ + --device_preset=ios-tof \ + --key_frame_distance=0.0001 \ + $base_dir/$scene + +frames=$(ls -l ${base_dir}/${scene}/rgb | wc -l) +if [ -z $frames ] || [ $frames -eq 0 ]; then + echo "ERROR: Not enough frames at $base_dir/$scene/rgb!" + exit +fi + +echo "Building SplaTAM expected directories..." +python3 scripts/nerfcapture2dataset.py --config $1 --base_dir $base_dir --scene $scene --frames $frames # Run SplaTAM -python3 scripts/splatam.py $1 +echo "Running SplaTAM..." +python3 scripts/splatam.py $1 --base_dir $base_dir --scene $scene --frames $frames # Visualize SplaTAM Output -python3 viz_scripts/final_recon.py $1 \ No newline at end of file +echo "Running visualizer..." +python3 viz_scripts/final_recon.py $1 --base_dir $base_dir --scene $scene --frames $frames + diff --git a/configs/iphone/nerfcapture.py b/configs/iphone/nerfcapture.py index f5b1c99d..09ac92b4 100644 --- a/configs/iphone/nerfcapture.py +++ b/configs/iphone/nerfcapture.py @@ -10,8 +10,8 @@ depth_scale = 10.0 # Depth Scale used when saving depth overwrite = False # Rewrite over dataset if it exists -full_res_width = 1920 -full_res_height = 1440 +full_res_width = 640 +full_res_height = 480 downscale_factor = 2.0 densify_downscale_factor = 4.0 @@ -145,4 +145,4 @@ viz_fps=5, # FPS for Online Recon Viz enter_interactive_post_online=False, # Enter Interactive Mode after Online Recon Viz ), -) \ No newline at end of file +) diff --git a/datasets/README.md b/datasets/README.md new file mode 100644 index 00000000..1e396836 --- /dev/null +++ b/datasets/README.md @@ -0,0 +1,4 @@ +# Datasets +## Requesting Ekumen HQ recorded datasets +If you wish to test out one of the datasets recorded at HQ please contact research@ekumenlabs.com. + diff --git a/docker/run.sh b/docker/run.sh index 665e7a35..7d33515c 100755 --- a/docker/run.sh +++ b/docker/run.sh @@ -1,11 +1,12 @@ #!/bin/bash timestamp=$(date +%s) +CONTAINER_NAME=splatam-$timestamp +IMAGE_NAME=splatam-env docker run -it \ --volume ./:/ws/SplaTAM/ \ --volume /tmp/.X11-unix:/tmp/.X11-unix \ - --rm \ --env NVIDIA_VISIBLE_DEVICES=all \ --env NVIDIA_DRIVER_CAPABILITIES=all \ --env DISPLAY=$DISPLAY \ @@ -15,9 +16,27 @@ docker run -it \ --group-add video \ --ulimit memlock=-1 \ --ulimit stack=67108864 \ - --name splatam-$timestamp \ + --name $CONTAINER_NAME \ --ipc=host \ --gpus all \ - splatam-env \ + $IMAGE_NAME \ /bin/bash +# Trap workspace exits and give the user the choice to save changes. +function onexit() { + while true; do + read -p "Do you want to overwrite the image called '$IMAGE_NAME' with the current changes? [y/n]: " answer + if [[ "${answer:0:1}" =~ y|Y ]]; then + echo "Overwriting docker image..." + docker commit $CONTAINER_NAME $IMAGE_NAME + break + elif [[ "${answer:0:1}" =~ n|N ]]; then + break + fi + done + docker stop $CONTAINER_NAME > /dev/null + docker rm $CONTAINER_NAME > /dev/null +} + +trap onexit EXIT + diff --git a/scripts/nerfcapture2dataset.py b/scripts/nerfcapture2dataset.py index 489ad95f..f754d183 100644 --- a/scripts/nerfcapture2dataset.py +++ b/scripts/nerfcapture2dataset.py @@ -5,6 +5,7 @@ #!/usr/bin/env python3 import argparse +import json import os import shutil import sys @@ -19,64 +20,17 @@ import cv2 import numpy as np -import cyclonedds.idl as idl -import cyclonedds.idl.annotations as annotate -import cyclonedds.idl.types as types -from dataclasses import dataclass -from cyclonedds.domain import DomainParticipant, Domain -from cyclonedds.core import Qos, Policy -from cyclonedds.sub import DataReader -from cyclonedds.topic import Topic -from cyclonedds.util import duration - def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--config", default="./configs/iphone/nerfcapture.py", type=str, help="Path to config file.") + parser.add_argument("--base_dir", type=str, help="Path to the dataset.") + parser.add_argument("--scene", type=str, help="Name of the NeRFCapture dataset. Usually has a _nerfcapture suffix.") + parser.add_argument("--frames", type=int, help="Amount of frames to process.") return parser.parse_args() -# DDS -# ================================================================================================== -@dataclass -@annotate.final -@annotate.autoid("sequential") -class SplatCaptureFrame(idl.IdlStruct, typename="SplatCaptureData.SplatCaptureFrame"): - id: types.uint32 - annotate.key("id") - timestamp: types.float64 - fl_x: types.float32 - fl_y: types.float32 - cx: types.float32 - cy: types.float32 - transform_matrix: types.array[types.float32, 16] - width: types.uint32 - height: types.uint32 - image: types.sequence[types.uint8] - has_depth: bool - depth_width: types.uint32 - depth_height: types.uint32 - depth_scale: types.float32 - depth_image: types.sequence[types.uint8] - - -dds_config = """ \ - \ - \ - \ - 10MB \ - \ - \ - config \ - stdout \ - \ - \ - \ -""" -# ================================================================================================== - - -def dataset_capture_loop(reader: DataReader, save_path: Path, overwrite: bool, n_frames: int, depth_scale: float): +def dataset_capture_loop(json_data, save_path: Path, overwrite: bool, n_frames: int, depth_scale: float): if save_path.exists(): if overwrite: # Prompt user to confirm deletion @@ -103,72 +57,68 @@ def dataset_capture_loop(reader: DataReader, save_path: Path, overwrite: bool, n total_frames = 0 # Total frames received - # Start DDS Loop - while True: - sample = reader.read_next() # Get frame from NeRFCapture - if sample: - print(f"{total_frames + 1}/{n_frames} frames received") - - if total_frames == 0: - save_path.mkdir(parents=True) - images_dir.mkdir() - manifest["w"] = sample.width - manifest["h"] = sample.height - manifest["cx"] = sample.cx - manifest["cy"] = sample.cy - manifest["fl_x"] = sample.fl_x - manifest["fl_y"] = sample.fl_y - manifest["integer_depth_scale"] = float(depth_scale)/65535.0 - if sample.has_depth: - depth_dir = save_path.joinpath("depth") - depth_dir.mkdir() - - # RGB - image = np.asarray(sample.image, dtype=np.uint8).reshape((sample.height, sample.width, 3)) - cv2.imwrite(str(images_dir.joinpath(f"{total_frames}.png")), cv2.cvtColor(image, cv2.COLOR_RGB2BGR)) - - # Depth if avaiable - depth = None - if sample.has_depth: - depth = np.asarray(sample.depth_image, dtype=np.uint8).view( - dtype=np.float32).reshape((sample.depth_height, sample.depth_width)) - depth = (depth*65535/float(depth_scale)).astype(np.uint16) - depth = cv2.resize(depth, dsize=( - sample.width, sample.height), interpolation=cv2.INTER_NEAREST) - cv2.imwrite(str(depth_dir.joinpath(f"{total_frames}.png")), depth) - - # Transform - X_WV = np.asarray(sample.transform_matrix, - dtype=np.float32).reshape((4, 4)).T - - frame = { - "transform_matrix": X_WV.tolist(), - "file_path": f"rgb/{total_frames}.png", - "fl_x": sample.fl_x, - "fl_y": sample.fl_y, - "cx": sample.cx, - "cy": sample.cy, - "w": sample.width, - "h": sample.height - } - - if depth is not None: - frame["depth_path"] = f"depth/{total_frames}.png" - - manifest["frames"].append(frame) - - # Update index - if total_frames == n_frames - 1: - print("Saving manifest...") - # Write manifest as json - manifest_json = json.dumps(manifest, indent=4) - with open(save_path.joinpath("transforms.json"), "w") as f: - f.write(manifest_json) - print("Done") - sys.exit(0) - total_frames += 1 + for sample in json_data['frames']: + has_depth = 'depth_path' in sample.keys() + print(f"{total_frames + 1}/{n_frames} frames received") + + if total_frames == 0: + save_path.mkdir(parents=True) + images_dir.mkdir() + manifest["w"] = sample.width + manifest["h"] = sample.height + manifest["cx"] = sample.cx + manifest["cy"] = sample.cy + manifest["fl_x"] = sample.fl_x + manifest["fl_y"] = sample.fl_y + manifest["integer_depth_scale"] = float(depth_scale)/65535.0 + if has_depth: + depth_dir = save_path.joinpath("depth") + depth_dir.mkdir() + + # RGB + image = cv2.imread(f"{sample['file_path']}.png") + cv2.imwrite(str(images_dir.joinpath(f"{total_frames}.png")), image) + + # Depth if avaiable + depth = None + if has_depth: + depth=cv2.imread(sample['depth_path']) + cv2.imwrite(str(depth_dir.joinpath(f"{total_frames}.png")), depth) + + # Transform + X_WV = np.asarray(sample.transform_matrix, + dtype=np.float32).reshape((4, 4)).T + + frame = { + "transform_matrix": X_WV.tolist(), + "file_path": f"rgb/{total_frames}.png", + "fl_x": sample.fl_x, + "fl_y": sample.fl_y, + "cx": sample.cx, + "cy": sample.cy, + "w": sample.width, + "h": sample.height + } + + if depth is not None: + frame["depth_path"] = f"depth/{total_frames}.png" + + manifest["frames"].append(frame) + + # Update index + if total_frames == n_frames - 1: + print("Saving manifest...") + # Write manifest as json + manifest_json = json.dumps(manifest, indent=4) + with open(save_path.joinpath("transforms.json"), "w") as f: + f.write(manifest_json) + print("Done") + sys.exit(0) + total_frames += 1 +from pathlib import Path + if __name__ == "__main__": args = parse_args() @@ -177,13 +127,15 @@ def dataset_capture_loop(reader: DataReader, save_path: Path, overwrite: bool, n os.path.basename(args.config), args.config ).load_module() - # Setup DDS - domain = Domain(domain_id=0, config=dds_config) - participant = DomainParticipant() - qos = Qos(Policy.Reliability.Reliable( - max_blocking_time=duration(seconds=1))) - topic = Topic(participant, "Frames", SplatCaptureFrame, qos=qos) - reader = DataReader(participant, topic) - + # TODO(Santoi): Improve config sharing between scripts. config = experiment.config - dataset_capture_loop(reader, Path(config['workdir']), config['overwrite'], config['num_frames'], config['depth_scale']) + config['workdir'] = args.base_dir + "/" + args.scene + config['data']['num_frames'] = config['num_frames'] = args.frames / 3 # SplaTAM usually stars to go nuts after 1/3 of the dataset frames. + config['data']['base_dir'] = config['base_dir'] = args.base_dir + config['data']['sequence'] = config['scene_name'] = args.scene + + transform_file = os.path.join(config['workdir'], 'transforms.json') + with open(transform_file, 'r') as f: + json_data = json.load(f) + + dataset_capture_loop(json_data, Path(config['workdir']), config['overwrite'], config['num_frames'], config['depth_scale']) diff --git a/scripts/splatam.py b/scripts/splatam.py index 1b082f7e..7edb8c8a 100644 --- a/scripts/splatam.py +++ b/scripts/splatam.py @@ -992,7 +992,10 @@ def rgbd_slam(config: dict): if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("experiment", type=str, help="Path to experiment file") + parser.add_argument("experiment", default="./configs/iphone/nerfcapture.py", type=str, help="Path to experiment file") + parser.add_argument("--base_dir", type=str, help="Path to the dataset.") + parser.add_argument("--scene", type=str, help="Name of the NeRFCapture dataset. Usually has a _nerfcapture suffix.") + parser.add_argument("--frames", type=int, help="Amount of frames to process.") args = parser.parse_args() @@ -1000,6 +1003,11 @@ def rgbd_slam(config: dict): os.path.basename(args.experiment), args.experiment ).load_module() + # TODO(Santoi): Improve config sharing between scripts. + experiment.config['workdir'] = args.base_dir + "/" + args.scene + experiment.config['data']['num_frames'] = experiment.config['num_frames'] = int(args.frames / 3) # SplaTAM usually stars to go nuts after 1/3 of the dataset frames. + experiment.config['data']['basedir'] = experiment.config['basedir'] = args.base_dir + experiment.config['data']['sequence'] = experiment.config['scene_name'] = args.scene # Set Experiment Seed seed_everything(seed=experiment.config['seed']) @@ -1011,4 +1019,4 @@ def rgbd_slam(config: dict): os.makedirs(results_dir, exist_ok=True) shutil.copy(args.experiment, os.path.join(results_dir, "config.py")) - rgbd_slam(experiment.config) \ No newline at end of file + rgbd_slam(experiment.config) diff --git a/spectacularAI-sdk b/spectacularAI-sdk new file mode 160000 index 00000000..b8c2ead8 --- /dev/null +++ b/spectacularAI-sdk @@ -0,0 +1 @@ +Subproject commit b8c2ead8f4905c4de9439b0689e3eef4fd53ddd4 diff --git a/viz_scripts/final_recon.py b/viz_scripts/final_recon.py index de1b7219..05a4bc34 100644 --- a/viz_scripts/final_recon.py +++ b/viz_scripts/final_recon.py @@ -279,6 +279,9 @@ def visualize(scene_path, cfg): parser = argparse.ArgumentParser() parser.add_argument("experiment", type=str, help="Path to experiment file") + parser.add_argument("--base_dir", type=str, help="Path to the dataset.") + parser.add_argument("--scene", type=str, help="Name of the NeRFCapture dataset. Usually has a _nerfcapture suffix.") + parser.add_argument("--frames", type=int, help="Amount of frames to process.") args = parser.parse_args() @@ -286,6 +289,11 @@ def visualize(scene_path, cfg): os.path.basename(args.experiment), args.experiment ).load_module() + # TODO(Santoi): Improve config sharing between scripts. + experiment.config['workdir'] = args.base_dir + "/" + args.scene + experiment.config['data']['num_frames'] = experiment.config['num_frames'] = int(args.frames / 3) # SplaTAM usually stars to go nuts after 1/3 of the dataset frames. + experiment.config['data']['basedir'] = experiment.config['basedir'] = args.base_dir + experiment.config['data']['sequence'] = experiment.config['scene_name'] = args.scene seed_everything(seed=experiment.config["seed"]) if "scene_path" not in experiment.config: