Initial commit for the code release of NeuRAD

Including: - dataparsers for several AD datasets - support for lidar data and moving actors - code for NeuRAD model - and more, check the code and readme :)
georghess · Apr 22, 2024 · ba30973 · ba30973
1 parent c572eb7
commit ba30973
Show file tree

Hide file tree

Showing 278 changed files with 10,448 additions and 451,246 deletions.
diff --git a/.gitignore b/.gitignore
@@ -165,6 +165,9 @@ cython_debug/
 outputs/
 exports/
 renders/
+*.mp4
+*.gif
+*.zip
 # tensorboard log files
 events.out.*
 
@@ -173,7 +176,8 @@ events.out.*
 
 # Misc
 old/
-temp*
+temp/
+tmp/
 !temporal*
 .nfs*
 external/
@@ -187,4 +191,5 @@ camera_paths/
 ._.DS_Store
 */**/.DS_Store
 */**/._.DS_Store
-
+masks/
+*.sif
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -3,14 +3,6 @@
 default_language_version:
     python: python3
 repos:
--   repo: local
-    hooks:
-    -   id: add-license-headers
-        name: add-license-headers
-        entry: nerfstudio/scripts/licensing/license_headers.sh
-        language: script
-        files: '.*'
-        pass_filenames: false
 -   repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.5.0
     hooks:
@@ -24,3 +16,10 @@ repos:
         args: [ --fix ]
     -   id: ruff-format
         types_or: [ python, pyi, jupyter ]
+-   repo: local
+    hooks:
+    -   id: check-copyright
+        name: Check Copyright
+        entry: python3 -m nerfstudio.scripts.licensing.check_copyright
+        language: python
+        types: [ python ]
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -27,127 +27,96 @@
       "justMyCode": true
     },
     {
-      "name": "Python: train",
+      "name": "Python: NeuRAD",
       "type": "python",
       "request": "launch",
       "program": "nerfstudio/scripts/train.py",
       "console": "integratedTerminal",
+      "justMyCode": true,
       "args": [
-        "instant_ngp",
-        "--viewer.no-enable",
-        "--logging.event-writer=wandb"
+        "neurad",
+        "--pipeline.model.eval_num_rays_per_chunk=2048",
+        "--pipeline.datamanager.num_processes=0",
+        "--pipeline.datamanager.train-num-lidar-rays-per-batch=1280",
+        "--pipeline.datamanager.eval-num-lidar-rays-per-batch=128",
+        "--pipeline.datamanager.train_num_rays_per_batch=4096",
+        "--pipeline.datamanager.eval_num_rays_per_batch=4096",
+        "pandaset-data",
+        "--data=./data/pandaset",
+        "--dataset_end_fraction=0.2",
+        "--cameras=front"
       ]
     },
     {
-      "name": "Python: train ngp",
-      "type": "python",
-      "request": "launch",
-      "program": "nerfstudio/scripts/train.py",
-      "console": "integratedTerminal",
-      "args": ["instant_ngp"]
-    },
-    {
-      "name": "Python: train ngp nerfstudio",
+      "name": "Python: NeuRAD tiny",
       "type": "python",
       "request": "launch",
       "program": "nerfstudio/scripts/train.py",
       "console": "integratedTerminal",
+      "justMyCode": true,
       "args": [
-        "instant_ngp",
-        // "--experiment_name=nerfstudio_experiment",
-        "--logging.local-writer.max-log-size=0",
-        "--viewer.no-enable"
-        // viewer
-        // "--viewer.enable",
-        // "--viewer.zmq-url=tcp://127.0.0.1:6001",
-        // "--viewer.websocket-port=8014",
-        // "--pipeline.model.randomize-background",
-        // "--pipeline.model.cone-angle=0.00390625",
-        // "nerfstudio-data"
+        "neurad",
+        "--pipeline.model.eval_num_rays_per_chunk=2048",
+        "--pipeline.model.field.grid.static.hashgrid_dim=2",
+        "--pipeline.model.field.grid.static.num_levels=4",
+        "--pipeline.model.field.grid.static.log2_hashmap_size=15",
+        "--pipeline.model.field.grid.actor.hashgrid_dim=2",
+        "--pipeline.model.field.grid.actor.num_levels=2",
+        "--pipeline.model.field.grid.actor.log2_hashmap_size=11",
+        "--pipeline.datamanager.num_processes=0",
+        "--pipeline.datamanager.train-num-lidar-rays-per-batch=1280",
+        "--pipeline.datamanager.eval-num-lidar-rays-per-batch=128",
+        "--pipeline.datamanager.train-num-rays-per_batch=2048",
+        "--pipeline.datamanager.eval-num-rays_per_batch=2048",
+        "pandaset-data",
+        "--data=./data/pandaset",
+        "--dataset_end_fraction=0.2",
+        "--cameras=front",
       ]
     },
     {
-      "name": "Python: train ngp-friends",
-      "type": "python",
-      "request": "launch",
-      "program": "nerfstudio/scripts/train.py",
-      "console": "integratedTerminal",
-      "args": [
-        "instant_ngp",
-        "--viewer.enable",
-        "--viewer.zmq-url=tcp://127.0.0.1:6001",
-        "--viewer.websocket-port=8014",
-        "--steps_per_save=1000000",
-        "--steps_per_test=1000000",
-        "--pipeline.model.density-field-params.base-scale=4",
-        "--pipeline.model.density-field-params.num-cascades=1",
-        "pipeline.datamanager.train-dataparser:friends-data-parser-config",
-        "--pipeline.datamanager.train-dataparser.scene-scale=4"
-      ]
-    },
-    {
-      "name": "Python: train semantic_nerfw",
-      "type": "python",
-      "request": "launch",
-      "program": "nerfstudio/scripts/train.py",
-      "console": "integratedTerminal",
-      "args": ["semantic_nerf"]
-    },
-    {
-      "name": "Python: NeuS on Replica",
+      "name": "Python: NeuRAD CPU",
       "type": "python",
       "request": "launch",
       "program": "nerfstudio/scripts/train.py",
       "console": "integratedTerminal",
       "justMyCode": false,
-      "env": { "CUDA_VISIBLE_DEVICES": "0" },
+      "env": {"PYTORCH_ENABLE_MPS_FALLBACK": "1"},
       "args": [
-        "neus",
-        "--pipeline.model.sdf-field.inside-outside=True",
-        "--pipeline.model.mono-depth-loss-mult=0.1",
-        "--pipeline.model.mono-normal-loss-mult=0.05",
-        "--vis=wandb",
-        "sdfstudio-data",
-        "--data=data/sdfstudio-demo-data/replica-room0",
-        "--include_mono_prior=True"
+        "neurad",
+        "--pipeline.model.eval_num_rays_per_chunk=2048",
+        "--pipeline.model.implementation=torch",
+        "--pipeline.model.sampling.proposal_field_1.grid.actor.use_4d_hashgrid=False",
+        "--pipeline.model.sampling.proposal_field_2.grid.actor.use_4d_hashgrid=False",
+        "--pipeline.model.field.grid.actor.use_4d_hashgrid=False",
+        "--machine.device-type=mps",
+        "--mixed_precision=False",
+        "--pipeline.datamanager.num_processes=0",
+        "--pipeline.datamanager.train-num-lidar-rays-per-batch=128",
+        "--pipeline.datamanager.eval-num-lidar-rays-per-batch=128",
+        "--pipeline.datamanager.train_num_rays_per_batch=2048",
+        "--pipeline.datamanager.eval_num_rays_per_batch=2048",
+        "--vis=tensorboard",
+        "pandaset-data",
+        "--data=./data/pandaset",
+        "--dataset_end_fraction=0.2",
+        "--cameras=front"
       ]
     },
     {
-      "name": "Python: NeuS-facto on Replica",
+      "name": "Python: Splatfacto PandaSet",
       "type": "python",
       "request": "launch",
       "program": "nerfstudio/scripts/train.py",
       "console": "integratedTerminal",
-      "justMyCode": false,
-      "env": { "CUDA_VISIBLE_DEVICES": "0" },
+      "justMyCode": true,
       "args": [
-        "neus-facto",
-        "--pipeline.model.sdf-field.inside-outside=True",
-        "--pipeline.model.mono-depth-loss-mult=0.1",
-        "--pipeline.model.mono-normal-loss-mult=0.05",
-        "--vis=wandb",
-        "sdfstudio-data",
-        "--data=data/sdfstudio-demo-data/replica-room0",
-        "--include_mono_prior=True"
+        "splatfacto",
+        "pandaset-data",
+        "--sequence=028",
+        "--data=./data/pandaset",
+        "--dataset_end_fraction=1.0",
       ]
     },
-    {
-      "name": "Python: NeuS-facto on DTU (no mono-prior)",
-      "type": "python",
-      "request": "launch",
-      "program": "nerfstudio/scripts/train.py",
-      "console": "integratedTerminal",
-      "justMyCode": false,
-      "env": { "CUDA_VISIBLE_DEVICES": "0" },
-      "args": [
-        "neus-facto",
-        "--pipeline.model.sdf-field.inside-outside=False",
-        "--vis=wandb",
-        "sdfstudio-data",
-        "--data=data/sdfstudio-demo-data/dtu-scan65",
-        "--include_mono_prior=False",
-        "--auto-orient=True"
-      ]
-    }
   ]
 }
diff --git a/Dockerfile b/Dockerfile
@@ -1,6 +1,5 @@
 ARG CUDA_VERSION=11.8.0
 ARG OS_VERSION=22.04
-ARG USER_ID=1000
 # Define base image.
 FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${OS_VERSION}
 ARG CUDA_VERSION
@@ -15,8 +14,8 @@ LABEL org.opencontainers.image.base.name="docker.io/library/nvidia/cuda:${CUDA_V
 
 # Variables used at build time.
 ## CUDA architectures, required by Colmap and tiny-cuda-nn.
-## NOTE: All commonly used GPU architectures are included and supported here. To speedup the image build process remove all architectures but the one of your explicit GPU. Find details here: https://developer.nvidia.com/cuda-gpus (8.6 translates to 86 in the line below) or in the docs.
-ARG CUDA_ARCHITECTURES=90;89;86;80;75;70;61;52;37
+## NOTE: Most commonly used GPU architectures are included and supported here. To speedup the image build process remove all architectures but the one of your explicit GPU. Find details here: https://developer.nvidia.com/cuda-gpus (8.6 translates to 86 in the line below) or in the docs.
+ARG CUDA_ARCHITECTURES=90;89;86;80;75
 
 # Set environment variables.
 ## Set non-interactive to prevent asking for user inputs blocking image creation.
@@ -53,18 +52,16 @@ RUN apt-get update && \
     libqt5opengl5-dev \
     libsqlite3-dev \
     libsuitesparse-dev \
-    nano \
     protobuf-compiler \
     python-is-python3 \
     python3.10-dev \
     python3-pip \
     qtbase5-dev \
-    sudo \
     vim-tiny \
-    wget && \
+    wget \
+    && \
     rm -rf /var/lib/apt/lists/*
 
-
 # Install GLOG (required by ceres).
 RUN git clone --branch v0.6.0 https://github.com/google/glog.git --single-branch && \
     cd glog && \
@@ -102,22 +99,6 @@ RUN git clone --branch 3.8 https://github.com/colmap/colmap.git --single-branch
     cd ../.. && \
     rm -rf colmap
 
-# Create non root user and setup environment.
-RUN useradd -m -d /home/user -g root -G sudo -u ${USER_ID} user
-RUN usermod -aG sudo user
-# Set user password
-RUN echo "user:user" | chpasswd
-# Ensure sudo group users are not asked for a password when using sudo command by ammending sudoers file
-RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
-
-# Switch to new uer and workdir.
-USER ${USER_ID}
-WORKDIR /home/user
-
-# Add local user binary folder to PATH variable.
-ENV PATH="${PATH}:/home/user/.local/bin"
-SHELL ["/bin/bash", "-c"]
-
 # Upgrade pip and install packages.
 RUN python3.10 -m pip install --no-cache-dir --upgrade pip setuptools pathtools promise pybind11
 # Install pytorch and submodules
@@ -158,15 +139,10 @@ RUN git clone --recursive https://github.com/cvg/pixel-perfect-sfm.git && \
 
 RUN python3.10 -m pip install --no-cache-dir omegaconf
 # Copy nerfstudio folder and give ownership to user.
-ADD . /home/user/nerfstudio
-USER root
-RUN chown -R user /home/user/nerfstudio
-USER ${USER_ID}
+ADD . /nerfstudio
 
 # Install nerfstudio dependencies.
-RUN cd nerfstudio && \
-    python3.10 -m pip install --no-cache-dir -e . && \
-    cd ..
+RUN cd /nerfstudio && python3.10 -m pip install --no-cache-dir -e .
 
 # Change working directory
 WORKDIR /workspace