diff --git a/README.md b/README.md index 8730cb20a..57d77724b 100644 --- a/README.md +++ b/README.md @@ -101,14 +101,14 @@ For CUDA 11.3: ```bash pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html -pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch +pip install ninja git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch ``` For CUDA 11.7: ```bash -pip install torch torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117 -pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch +pip install torch==1.13.1 torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117 +pip install ninja git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch ``` See [Dependencies](https://github.com/nerfstudio-project/nerfstudio/blob/main/docs/quickstart/installation.md#dependencies) @@ -196,16 +196,17 @@ ns-export pointcloud --help Using an existing dataset is great, but likely you want to use your own data! We support various methods for using your own data. Before it can be used in nerfstudio, the camera location and orientations must be determined and then converted into our format using `ns-process-data`. We rely on external tools for this, instructions and information can be found in the documentation. -| Data | Capture Device | Requirements | `ns-process-data` Speed | -| ------------------------------------------------------------------------------------------------- | -------------- | ----------------------------------------------- | ----------------------- | -| 📷 [Images](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#images-and-video) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | -| 📹 [Video](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#images-and-video) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | -| 🌎 [360 Data](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#360_data) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | -| 📱 [Polycam](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#polycam-capture) | IOS with LiDAR | [Polycam App](https://poly.cam/) | 🐇 | -| 📱 [KIRI Engine](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#kiri-capture) | IOS or Android | [KIRI Engine App](https://www.kiriengine.com/) | 🐇 | -| 📱 [Record3D](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#record3d-capture) | IOS with LiDAR | [Record3D app](https://record3d.app/) | 🐇 | -| 🖥 [Metashape](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#metashape) | Any | [Metashape](https://www.agisoft.com/) | 🐇 | -| 🛠 [Custom](https://docs.nerf.studio/en/latest/quickstart/data_conventions.html) | Any | Camera Poses | 🐇 | +| Data | Capture Device | Requirements | `ns-process-data` Speed | +| ---------------------------------------------------------------------------------------------------- | -------------- | ----------------------------------------------------------------- | ----------------------- | +| 📷 [Images](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#images-and-video) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | +| 📹 [Video](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#images-and-video) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | +| 🌎 [360 Data](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#360_data) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | +| 📱 [Polycam](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#polycam-capture) | IOS with LiDAR | [Polycam App](https://poly.cam/) | 🐇 | +| 📱 [KIRI Engine](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#kiri-capture) | IOS or Android | [KIRI Engine App](https://www.kiriengine.com/) | 🐇 | +| 📱 [Record3D](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#record3d-capture) | IOS with LiDAR | [Record3D app](https://record3d.app/) | 🐇 | +| 🖥 [Metashape](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#metashape) | Any | [Metashape](https://www.agisoft.com/) | 🐇 | +| 🖥 [RealityCapture](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#realitycapture) | Any | [RealityCapture](https://www.capturingreality.com/realitycapture) | 🐇 | +| 🛠 [Custom](https://docs.nerf.studio/en/latest/quickstart/data_conventions.html) | Any | Camera Poses | 🐇 | ## 5. Advanced Options diff --git a/colab/demo.ipynb b/colab/demo.ipynb index 4a87f1718..541b9d741 100644 --- a/colab/demo.ipynb +++ b/colab/demo.ipynb @@ -7,7 +7,7 @@ "colab_type": "text" }, "source": [ - "\"Open" + "\"Open" ] }, { @@ -329,7 +329,7 @@ "#@markdown
The rendered video should be at renders/output.mp4!
\n", "\n", "\n", - "base_dir = \"/content/outputs/data-nerfstudio-\" + scene + \"/nerfacto/\"\n", + "base_dir = \"/content/outputs/unnamed/nerfacto/\"\n", "training_run_dir = base_dir + os.listdir(base_dir)[0]\n", "\n", "from IPython.core.display import display, HTML\n", diff --git a/docs/developer_guides/config.md b/docs/developer_guides/config.md index 29d093f9d..bf216278f 100644 --- a/docs/developer_guides/config.md +++ b/docs/developer_guides/config.md @@ -118,33 +118,3 @@ Often times, you just want to play with the parameters of an existing model with # NOTE: the dataparser and associated configurations go at the end of the command ns-train {METHOD_NAME} --vis viewer {DATA_PARSER} --scale-factor 0.5 ``` - -### Extending Nerfstudio with custom methods -In order to extend the Nerfstudio and register your own methods, you can package your code as a python package -and register it with Nerfstudio as a `nerfstudio.method_configs` entrypoint in the `pyproject.toml` file. -The Nerfstudio will automatically look for all registered methods and will register them to be used -by method such as `ns-train`. - -Here is an example: -```python -"""my_package/my_config.py""" - -from nerfstudio.engine.trainer import TrainerConfig -from nerfstudio.plugins.types import MethodSpecification - -MyMethod = MethodSpecification( - config=TrainerConfig( - method_name="my-method", - ... - ), - description="Custom description" -) - -"""pyproject.toml""" -[project] -name = "my_package" -... - -[project.entry-points.'nerfstudio.method_configs'] -my-method = 'my_package.my_config:MyMethod' -``` \ No newline at end of file diff --git a/docs/developer_guides/new_methods.md b/docs/developer_guides/new_methods.md new file mode 100644 index 000000000..c0484ca9e --- /dev/null +++ b/docs/developer_guides/new_methods.md @@ -0,0 +1,75 @@ +# Adding a New Method + +Nerfstudio aims to offer researchers a codebase that they can utilize to extend and develop novel methods. Our vision is for users to establish a distinct repository that imports nerfstudio and overrides pipeline components to cater to specific functionality requirements of the new approach. If any of the new features require modifications to the core nerfstudio repository and can be generally useful, we encourage you to submit a PR to enable others to benefit from it. + +Examples are often the best way to learn, take a look at the [LERF](https://github.com/kerrj/lerf) repository for good example of how to use nerfstudio in your projects. + +## File Structure + +We recommend the following file structure: + +``` +├── my_method +│ ├── my_config.py +│ ├── custom_pipeline.py [optional] +│ ├── custom_model.py [optional] +│ ├── custom_field.py [optional] +│ ├── custom_datamanger.py [optional] +│ ├── ... +├── pyproject.toml +``` + +## Registering custom model with nerfstudio + +In order to extend the Nerfstudio and register your own methods, you can package your code as a python package +and register it with Nerfstudio as a `nerfstudio.method_configs` entrypoint in the `pyproject.toml` file. +Nerfstudio will automatically look for all registered methods and will register them to be used +by methods such as `ns-train`. + +First create a config file: + +```python +"""my_method/my_config.py""" + +from nerfstudio.engine.trainer import TrainerConfig +from nerfstudio.plugins.types import MethodSpecification + +MyMethod = MethodSpecification( + config=TrainerConfig( + method_name="my-method", + pipeline=... + ... + ), + description="Custom description" +) +``` + +Then create a `pyproject.toml` file. This is where the entrypoint to your method is set and also where you can specify additional dependencies required by your codebase. + +```python +"""pyproject.toml""" + +[project] +name = "my_method" + +dependencies = [ + "nerfstudio" # you may want to consider pinning the version, ie "nerfstudio==0.1.19" +] + +[project.entry-points.'nerfstudio.method_configs'] +my-method = 'my_method.my_config:MyMethod' +``` + +finally run the following to register the method, + +``` +pip install -e . +``` + +## Running custom method + +After registering your method you should be able to run the method with, + +``` +ns-train my-method --data DATA_DIR +``` diff --git a/docs/extensions/imgs/sdfstudio_overview.svg b/docs/extensions/imgs/sdfstudio_overview.svg new file mode 100644 index 000000000..9183ef453 --- /dev/null +++ b/docs/extensions/imgs/sdfstudio_overview.svg @@ -0,0 +1,33888 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + ... + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/extensions/sdfstudio.md b/docs/extensions/sdfstudio.md new file mode 100644 index 000000000..8f8134662 --- /dev/null +++ b/docs/extensions/sdfstudio.md @@ -0,0 +1,50 @@ +# SDFStudio + +[project website](https://autonomousvision.github.io/sdfstudio/) + +```{image} imgs/sdfstudio_overview.svg +:width: 800 +:align: center +:alt: sdfstudio overview figure +``` + +## Overview + +SDFStudio is built on top of nerfstudio. It implements multiple implicit surface reconstruction methods including: + +- UniSurf +- VolSDF +- NeuS +- MonoSDF +- Mono-UniSurf +- Mono-NeuS +- Geo-NeuS +- Geo-UniSurf +- Geo-VolSDF +- NeuS-acc +- NeuS-facto +- NeuralReconW + +You can learn more about these methods [here](https://github.com/autonomousvision/sdfstudio/blob/master/docs/sdfstudio-methods.md#Methods) + +## Surface models in nerfstudio + +We intend to integrate many of the SDFStudio improvments back into the nerstudio core repository. + +Supported methods: + +- NeuS + +## Citation + +If you use these surface based models in your research, you should consider citing the authors of SDFStudio, + +```none +@misc{Yu2022SDFStudio, + author = {Yu, Zehao and Chen, Anpei and Antic, Bozidar and Peng, Songyou Peng and Bhattacharyya, Apratim + and Niemeyer, Michael and Tang, Siyu and Sattler, Torsten and Geiger, Andreas}, + title = {SDFStudio: A Unified Framework for Surface Reconstruction}, + year = {2022}, + url = {https://github.com/autonomousvision/sdfstudio}, +} +``` diff --git a/docs/index.md b/docs/index.md index b03e8712b..d28393883 100644 --- a/docs/index.md +++ b/docs/index.md @@ -60,6 +60,7 @@ Contributing :caption: Extensions extensions/blender_addon extensions/unreal_engine +extensions/sdfstudio ``` ```{toctree} @@ -74,6 +75,7 @@ nerfology/model_components/index :hidden: :caption: Developer Guides +developer_guides/new_methods developer_guides/pipelines/index developer_guides/viewer/viewer_overview developer_guides/config diff --git a/docs/quickstart/custom_dataset.md b/docs/quickstart/custom_dataset.md index 4295bf4bc..65e6c22d4 100644 --- a/docs/quickstart/custom_dataset.md +++ b/docs/quickstart/custom_dataset.md @@ -5,7 +5,7 @@ Training model on existing datasets is only so fun. If you would like to train o To process your own data run: ```bash -ns-process-data {video,images,polycam,insta360,record3d} --data {DATA_PATH} --output-dir {PROCESSED_DATA_DIR} +ns-process-data {video,images,polycam,record3d} --data {DATA_PATH} --output-dir {PROCESSED_DATA_DIR} ``` A full set of arguments can be found {doc}`here`. @@ -20,6 +20,7 @@ We Currently support the following custom data types: | 📱 [KIRI Engine](kiri) | IOS or Android | [KIRI Engine App](https://www.kiriengine.com/) | 🐇 | | 📱 [Record3D](record3d) | IOS with LiDAR | [Record3D app](https://record3d.app/) | 🐇 | | 🖥 [Metashape](metashape) | Any | [Metashape](https://www.agisoft.com/) | 🐇 | +| 🖥 [RealityCapture](realitycapture) | Any | [RealityCapture](https://www.capturingreality.com/realitycapture) | 🐇 | (images_and_video)= @@ -33,12 +34,6 @@ To assist running on custom data we have a script that will process a video or f - COLMAP can be finicky. Try your best to capture overlapping, non-blurry images. ::: -:::{admonition} Note -:class: warning - -- If COLMAP opens in a new window, simply close it to resume processing. - ::: - ### Processing Data ```bash @@ -292,6 +287,26 @@ ns-process-data metashape --data {data directory} --xml {xml file} --output-dir ns-train nerfacto --data {output directory} ``` +(realitycapture)= + +## RealityCapture + +1. Align your images using RealityCapture. `ALIGNMENT -> Align Images` + +2. Export the camera alignment as a `csv` file. Choose `Internal/External camera parameters` + +3. Convert the data to the nerfstudio format. + +```bash +ns-process-data realitycapture --data {data directory} --csv {csv file} --output-dir {output directory} +``` + +4. Train with nerfstudio! + +```bash +ns-train nerfacto --data {output directory} +``` + (360_data)= ## 360 Data (Equirectangular) @@ -303,10 +318,10 @@ Equirectangular data is data that has been taken by a 360 camera such as Insta36 For a set of equirectangular images, process the data using the following command: ```bash -ns-process-data images --camera-type equirectangular --images-per-equirect {8, or 14} --data {data directory} --output-dir {output directory} +ns-process-data images --camera-type equirectangular --images-per-equirect {8, or 14} --crop-factor {top bottom left right} --data {data directory} --output-dir {output directory} ``` -The images-per-equirect argument is the number of images that will be sampled from each equirectangular image. We have found that 8 images per equirectangular image is sufficient for most use cases so it defaults to that. However, if you find that there isn't enough detail in the nerf or that colmap is having trouble aligning the images, you can try increasing the number of images per equirectangular image to 14. +The images-per-equirect argument is the number of images that will be sampled from each equirectangular image. We have found that 8 images per equirectangular image is sufficient for most use cases so it defaults to that. However, if you find that there isn't enough detail in the nerf or that colmap is having trouble aligning the images, you can try increasing the number of images per equirectangular image to 14. See the video section below for details on cropping. ### Videos diff --git a/docs/quickstart/installation.md b/docs/quickstart/installation.md index 950aefb2a..dafb9f0ee 100644 --- a/docs/quickstart/installation.md +++ b/docs/quickstart/installation.md @@ -71,7 +71,7 @@ pip uninstall torch torchvision functorch Install pytorch 1.13.1 with CUDA and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn) ```bash -pip install torch torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117 +pip install torch==1.13.1 torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117 ``` ::: @@ -79,10 +79,10 @@ pip install torch torchvision functorch --extra-index-url https://download.pytor ### tinycudann -After pytorch, install the torch bindings for tinycudann: +After pytorch and ninja, install the torch bindings for tinycudann: ```bash -pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch +pip install ninja git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch ``` ## Installing nerfstudio diff --git a/docs/reference/cli/index.md b/docs/reference/cli/index.md index 9e1b8fae4..7bc95834f 100644 --- a/docs/reference/cli/index.md +++ b/docs/reference/cli/index.md @@ -6,21 +6,22 @@ We provide a command line interface for training your own NeRFs (no coding neces Here are the popular commands that we offer. If you've cloned the repo, you can also look at the [pyproject.toml file](https://github.com/nerfstudio-project/nerfstudio/blob/main/pyproject.toml) at the `[project.scripts]` section for details. -| Command | Description | Filename | -| ---------------- | -------------------------------------- | ---------------------------------- | -| ns-install-cli | Install tab completion for all scripts | scripts/completions/install.py | -| ns-process-data | Generate a dataset from your own data | scripts/process_data.py | -| ns-download-data | Download existing captures | scripts/downloads/download_data.py | -| ns-train | Generate a NeRF | scripts/train.py | -| ns-viewer | View a trained NeRF | scripts/viewer/run_viewer.py | -| ns-eval | Run evaluation metrics for your Model | scripts/eval.py | -| ns-render | Render out a video of your NeRF | scripts/render.py | -| ns-export | Export a NeRF into other formats | scripts/exporter.py | +| Command | Description | Filename | +| ------------------------------------ | -------------------------------------- | ---------------------------------- | +| [ns-install-cli](ns_install_cli) | Install tab completion for all scripts | scripts/completions/install.py | +| [ns-process-data](ns_process_data) | Generate a dataset from your own data | scripts/process_data.py | +| [ns-download-data](ns_download_data) | Download existing captures | scripts/downloads/download_data.py | +| [ns-train](ns_train) | Generate a NeRF | scripts/train.py | +| [ns-viewer](ns_viewer) | View a trained NeRF | scripts/viewer/run_viewer.py | +| [ns-eval](ns_eval) | Run evaluation metrics for your Model | scripts/eval.py | +| [ns-render](ns_render) | Render out a video of your NeRF | scripts/render.py | +| [ns-export](ns_export) | Export a NeRF into other formats | scripts/exporter.py | ```{toctree} :maxdepth: 1 :hidden: +ns_install_cli ns_process_data ns_download_data ns_train diff --git a/docs/reference/cli/ns_install_cli.md b/docs/reference/cli/ns_install_cli.md new file mode 100644 index 000000000..e9c1a9541 --- /dev/null +++ b/docs/reference/cli/ns_install_cli.md @@ -0,0 +1,9 @@ +# ns-install-cli + +```{eval-rst} +.. argparse:: + :module: scripts.completions.install + :func: get_parser_fn + :prog: ns-install-cli + :nodefault: +``` diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index 9d95cc269..3d7872d74 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -23,11 +23,14 @@ from dataclasses import dataclass from pathlib import Path from typing import Any, List, Optional, Tuple, Type +from typing_extensions import Literal + # model instances from nerfstudio.utils import writer warnings.filterwarnings("ignore", module="torchvision") + # Pretty printing class class PrintableConfig: # pylint: disable=too-few-public-methods """Printable Config defining str function""" @@ -135,8 +138,10 @@ class ViewerConfig(PrintableConfig): """The zmq port to connect to for communication. If None, find an available port.""" launch_bridge_server: bool = True """whether or not to launch the bridge server""" - websocket_port: Optional[int] = 7007 - """the default websocket port to connect to""" + websocket_port: Optional[int] = None + """The websocket port to connect to. If None, find an available port.""" + websocket_port_default: int = 7007 + """The default websocket port to connect to if websocket_port is not specified""" ip_address: str = "127.0.0.1" """the ip address where the bridge server is running""" num_rays_per_chunk: int = 32768 @@ -146,3 +151,9 @@ class ViewerConfig(PrintableConfig): actually used in training/evaluation. If -1, display all.""" quit_on_train_completion: bool = False """Whether to kill the training job when it has completed. Note this will stop rendering in the viewer.""" + image_format: Literal["jpeg", "png"] = "jpeg" + """Image format viewer should use; jpeg is lossy compression, while png is lossless.""" + jpeg_quality: int = 90 + """Quality tradeoff to use for jpeg compression.""" + png_compression: int = 1 + """Size/speed tradeoff to use for png compression.""" diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index b80860f59..3e43406f2 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -104,7 +104,9 @@ method_configs["nerfacto"] = TrainerConfig( train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, camera_optimizer=CameraOptimizerConfig( - mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) + mode="SO3xR3", + optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2), + scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000), ), ), model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), @@ -112,11 +114,11 @@ method_configs["nerfacto"] = TrainerConfig( optimizers={ "proposal_networks": { "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, + "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), }, "fields": { "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, + "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), @@ -319,8 +321,8 @@ method_configs["tensorf"] = TrainerConfig( pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=BlenderDataParserConfig(), - train_num_rays_per_batch=2048, - eval_num_rays_per_batch=2048, + train_num_rays_per_batch=4096, + eval_num_rays_per_batch=4096, ), model=TensoRFModelConfig(), ), @@ -367,13 +369,18 @@ method_configs["phototourism"] = TrainerConfig( max_num_iterations=30000, mixed_precision=True, pipeline=VanillaPipelineConfig( - datamanager=VariableResDataManagerConfig( # NOTE: one of the only differences with nerfacto + datamanager=VanillaDataManagerConfig( dataparser=PhototourismDataParserConfig(), # NOTE: one of the only differences with nerfacto train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, camera_optimizer=CameraOptimizerConfig( mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) ), + # Large dataset, so using prior values from VariableResDataManager. + train_num_images_to_sample_from=40, + train_num_times_to_repeat_images=100, + eval_num_images_to_sample_from=40, + eval_num_times_to_repeat_images=100, ), model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), ), diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 7ff2f9aaf..8bcd21bbe 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -47,6 +47,7 @@ from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( InstantNGPDataParserConfig, ) from nerfstudio.data.dataparsers.minimal_dataparser import MinimalDataParserConfig +from nerfstudio.data.dataparsers.nerfosr_dataparser import NeRFOSRDataParserConfig from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig from nerfstudio.data.dataparsers.phototourism_dataparser import ( @@ -73,6 +74,31 @@ from nerfstudio.utils.misc import IterableWrapper CONSOLE = Console(width=120) + +def variable_res_collate(batch: List[Dict]) -> Dict: + """Default collate function for the cached dataloader. + Args: + batch: Batch of samples from the dataset. + Returns: + Collated batch. + """ + images = [] + masks = [] + for data in batch: + image = data.pop("image") + mask = data.pop("mask", None) + images.append(image) + if mask: + masks.append(mask) + + new_batch: dict = nerfstudio_collate(batch) + new_batch["image"] = images + if masks: + new_batch["mask"] = masks + + return new_batch + + AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of flags in subcommands. tyro.extras.subcommand_type_from_defaults( { @@ -87,6 +113,7 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of "dycheck-data": DycheckDataParserConfig(), "scannet-data": ScanNetDataParserConfig(), "sdfstudio-data": SDFStudioDataParserConfig(), + "nerfosr-data": NeRFOSRDataParserConfig(), "sitcoms3d-data": Sitcoms3DDataParserConfig(), }, prefix_names=False, # Omit prefixes in subcommands themselves. @@ -383,6 +410,16 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method self.train_dataset = self.create_train_dataset() self.eval_dataset = self.create_eval_dataset() + + if self.train_dataparser_outputs is not None: + cameras = self.train_dataparser_outputs.cameras + if len(cameras) > 1: + for i in range(1, len(cameras)): + if cameras[0].width != cameras[i].width or cameras[0].height != cameras[i].height: + CONSOLE.print("Variable resolution, using variable_res_collate") + self.config.collate_fn = variable_res_collate + break + super().__init__() def create_train_dataset(self) -> InputDataset: diff --git a/nerfstudio/data/datamanagers/variable_res_datamanager.py b/nerfstudio/data/datamanagers/variable_res_datamanager.py deleted file mode 100644 index b821bc034..000000000 --- a/nerfstudio/data/datamanagers/variable_res_datamanager.py +++ /dev/null @@ -1,66 +0,0 @@ -# Copyright 2022 The Nerfstudio Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Data loader for variable resolution datasets, where batching raw image tensors isn't possible. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Dict, List - -from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig -from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate - - -def variable_res_collate(batch: List[Dict]) -> Dict: - """Default collate function for the cached dataloader. - Args: - batch: Batch of samples from the dataset. - Returns: - Collated batch. - """ - images = [] - masks = [] - for data in batch: - image = data.pop("image") - mask = data.pop("mask", None) - images.append(image) - if mask: - masks.append(mask) - - new_batch: dict = nerfstudio_collate(batch) - new_batch["image"] = images - if masks: - new_batch["mask"] = masks - - return new_batch - - -@dataclass -class VariableResDataManagerConfig(VanillaDataManagerConfig): - """A datamanager for variable resolution datasets, with presets to optimize - for the fact that we are now dealing with lists of images and masks. - """ - - train_num_images_to_sample_from: int = 40 - """Number of images to sample during training iteration.""" - train_num_times_to_repeat_images: int = 100 - """When not training on all images, number of iterations before picking new - images. If -1, never pick new images.""" - eval_num_images_to_sample_from: int = 40 - """Number of images to sample during eval iteration.""" - eval_num_times_to_repeat_images: int = 100 - collate_fn = staticmethod(variable_res_collate) diff --git a/nerfstudio/data/dataparsers/base_dataparser.py b/nerfstudio/data/dataparsers/base_dataparser.py index 89c7c4e21..12367e51f 100644 --- a/nerfstudio/data/dataparsers/base_dataparser.py +++ b/nerfstudio/data/dataparsers/base_dataparser.py @@ -24,6 +24,7 @@ from typing import Any, Dict, List, Optional, Type import torch from torchtyping import TensorType +from typing_extensions import Literal import nerfstudio.configs.base_config as cfg from nerfstudio.cameras.cameras import Cameras @@ -89,6 +90,26 @@ class DataparserOutputs: with open(path, "w", encoding="UTF-8") as file: json.dump(data, file, indent=4) + def transform_poses_to_original_space( + self, + poses: TensorType["num_poses", 3, 4], + camera_convention: Literal["opengl", "opencv"] = "opencv", + ) -> TensorType["num_poses", 3, 4]: + """ + Transforms the poses in the transformed space back to the original world coordinate system. + Args: + poses: Poses in the transformed space + camera_convention: Camera system convention used for the transformed poses + Returns: + Original poses + """ + return transform_poses_to_original_space( + poses, + self.dataparser_transform, + self.dataparser_scale, + camera_convention=camera_convention, + ) + @dataclass class DataParserConfig(cfg.InstantiateConfig): @@ -136,3 +157,46 @@ class DataParser: """ dataparser_outputs = self._generate_dataparser_outputs(split) return dataparser_outputs + + +def transform_poses_to_original_space( + poses: TensorType["num_poses", 3, 4], + applied_transform: TensorType[3, 4], + applied_scale: float, + camera_convention: Literal["opengl", "opencv"] = "opencv", +) -> TensorType["num_poses", 3, 4]: + """ + Transforms the poses in the transformed space back to the original world coordinate system. + Args: + poses: Poses in the transformed space + applied_transform: Transform matrix applied in the data processing step + applied_scale: Scale used in the data processing step + camera_convention: Camera system convention used for the transformed poses + Returns: + Original poses + """ + output_poses = torch.cat( + ( + poses, + torch.tensor([[[0, 0, 0, 1]]], dtype=poses.dtype, device=poses.device).repeat_interleave(len(poses), 0), + ), + 1, + ) + output_poses[..., :3, 3] /= applied_scale + inv_transform = torch.linalg.inv( + torch.cat( + ( + applied_transform, + torch.tensor([[0, 0, 0, 1]], dtype=applied_transform.dtype, device=applied_transform.device), + ), + 0, + ) + ) + output_poses = torch.einsum("ij,bjk->bik", inv_transform, output_poses) + if camera_convention == "opencv": + output_poses[..., 0:3, 1:3] *= -1 + elif camera_convention == "opengl": + pass + else: + raise ValueError(f"Camera convention {camera_convention} is not supported.") + return output_poses[:, :3] diff --git a/nerfstudio/data/dataparsers/nerfosr_dataparser.py b/nerfstudio/data/dataparsers/nerfosr_dataparser.py new file mode 100644 index 000000000..820a28956 --- /dev/null +++ b/nerfstudio/data/dataparsers/nerfosr_dataparser.py @@ -0,0 +1,236 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" Data parser for NeRF-OSR datasets + + Presented in the paper: https://4dqv.mpi-inf.mpg.de/NeRF-OSR/ + +""" + +from __future__ import annotations + +import glob +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import List, Tuple, Type + +import numpy as np +import torch +from rich.console import Console +from typing_extensions import Literal + +from nerfstudio.cameras import camera_utils +from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.data.dataparsers.base_dataparser import ( + DataParser, + DataParserConfig, + DataparserOutputs, +) +from nerfstudio.data.scene_box import SceneBox + +CONSOLE = Console(width=120) + + +def find_files(directory: str, exts: List[str]): + """Find all files in a directory that have a certain file extension. + + Args: + directory : The directory to search for files. + exts : A list of file extensions to search for. Each file extension should be in the form '*.ext'. + + Returns: + A list of file paths for all the files that were found. The list is sorted alphabetically. + """ + if os.path.isdir(directory): + # types should be ['*.png', '*.jpg', '*.JPG', '*.PNG'] + files_grabbed = [] + for ext in exts: + files_grabbed.extend(glob.glob(os.path.join(directory, ext))) + if len(files_grabbed) > 0: + files_grabbed = sorted(files_grabbed) + return files_grabbed + return [] + + +def _parse_osm_txt(filename: str): + """Parse a text file containing numbers and return a 4x4 numpy array of float32 values. + + Args: + filename : a file containing numbers in a 4x4 matrix. + + Returns: + A numpy array of shape [4, 4] containing the numbers from the file. + """ + assert os.path.isfile(filename) + with open(filename, encoding="UTF-8") as f: + nums = f.read().split() + return np.array([float(x) for x in nums]).reshape([4, 4]).astype(np.float32) + + +def get_camera_params( + scene_dir: str, split: Literal["train", "validation", "test"] +) -> Tuple[torch.Tensor, torch.Tensor, int]: + """Load camera intrinsic and extrinsic parameters for a given scene split. + + Args" + scene_dir : The directory containing the scene data. + split : The split for which to load the camera parameters. + + Returns + A tuple containing the intrinsic parameters (as a torch.Tensor of shape [N, 4, 4]), + the camera-to-world matrices (as a torch.Tensor of shape [N, 4, 4]), and the number of cameras (N). + """ + split_dir = f"{scene_dir}/{split}" + + # camera parameters files + intrinsics_files = find_files(f"{split_dir}/intrinsics", exts=["*.txt"]) + pose_files = find_files(f"{split_dir}/pose", exts=["*.txt"]) + + num_cams = len(pose_files) + + intrinsics = [] + camera_to_worlds = [] + for i in range(num_cams): + intrinsics.append(_parse_osm_txt(intrinsics_files[i])) + + pose = _parse_osm_txt(pose_files[i]) + + # convert from COLMAP/OpenCV to nerfstudio camera (OpenGL/Blender) + pose[0:3, 1:3] *= -1 + + camera_to_worlds.append(pose) + + intrinsics = torch.from_numpy(np.stack(intrinsics).astype(np.float32)) # [N, 4, 4] + camera_to_worlds = torch.from_numpy(np.stack(camera_to_worlds).astype(np.float32)) # [N, 4, 4] + + return intrinsics, camera_to_worlds, num_cams + + +@dataclass +class NeRFOSRDataParserConfig(DataParserConfig): + """Nerfstudio dataset config""" + + _target: Type = field(default_factory=lambda: NeRFOSR) + """target class to instantiate""" + data: Path = Path("data/NeRF-OSR/Data/") + """Directory specifying location of data.""" + scene: str = "stjacob" + """Which scene to load""" + scene_scale: float = 1.0 + """How much to scale the region of interest by.""" + scale_factor: float = 1.0 + """How much to scale the camera origins by.""" + use_masks: bool = False + """Whether to use masks.""" + orientation_method: Literal["pca", "up", "vertical", "none"] = "vertical" + """The method to use for orientation.""" + center_method: Literal["poses", "focus", "none"] = "focus" + """The method to use for centering.""" + auto_scale_poses: bool = True + """Whether to automatically scale the poses to fit in +/- 1 bounding box.""" + + +@dataclass +class NeRFOSR(DataParser): + """NeRFOSR Dataparser + Presented in the paper: https://4dqv.mpi-inf.mpg.de/NeRF-OSR/ + + Some of this code comes from https://github.com/r00tman/NeRF-OSR/blob/main/data_loader_split.py + + Source data convention is: + camera coordinate system: x-->right, y-->down, z-->scene (opencv/colmap convention) + poses is camera-to-world + masks are 0 for dynamic content, 255 for static content + """ + + config: NeRFOSRDataParserConfig + + def _generate_dataparser_outputs(self, split="train"): + data = self.config.data + scene = self.config.scene + split = "validation" if split == "val" else split + + if scene == "trevi": + scene_dir = f"{data}/{scene}/final_clean" + split_dir = f"{data}/{scene}/final_clean/{split}" + else: + scene_dir = f"{data}/{scene}/final" + split_dir = f"{data}/{scene}/final/{split}" + + # get all split cam params + intrinsics_train, camera_to_worlds_train, n_train = get_camera_params(scene_dir, "train") + intrinsics_val, camera_to_worlds_val, n_val = get_camera_params(scene_dir, "validation") + intrinsics_test, camera_to_worlds_test, _ = get_camera_params(scene_dir, "test") + + # combine all cam params + intrinsics = torch.cat([intrinsics_train, intrinsics_val, intrinsics_test], dim=0) + camera_to_worlds = torch.cat([camera_to_worlds_train, camera_to_worlds_val, camera_to_worlds_test], dim=0) + + camera_to_worlds, _ = camera_utils.auto_orient_and_center_poses( + camera_to_worlds, + method=self.config.orientation_method, + center_method=self.config.center_method, + ) + + # Scale poses + scale_factor = 1.0 + if self.config.auto_scale_poses: + scale_factor /= torch.max(torch.abs(camera_to_worlds[:, :3, 3])) + + camera_to_worlds[:, :3, 3] *= scale_factor * self.config.scale_factor + + if split == "train": + camera_to_worlds = camera_to_worlds[:n_train] + intrinsics = intrinsics[:n_train] + elif split == "validation": + camera_to_worlds = camera_to_worlds[n_train : n_train + n_val] + intrinsics = intrinsics[n_train : n_train + n_val] + elif split == "test": + camera_to_worlds = camera_to_worlds[n_train + n_val :] + intrinsics = intrinsics[n_train + n_val :] + + cameras = Cameras( + camera_to_worlds=camera_to_worlds[:, :3, :4], + fx=intrinsics[:, 0, 0], + fy=intrinsics[:, 1, 1], + cx=intrinsics[:, 0, 2], + cy=intrinsics[:, 1, 2], + camera_type=CameraType.PERSPECTIVE, + ) + + # in x,y,z order + # assumes that the scene is centered at the origin + aabb_scale = self.config.scene_scale + scene_box = SceneBox( + aabb=torch.tensor( + [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 + ) + ) + + # --- images --- + image_filenames = find_files(f"{split_dir}/rgb", exts=["*.png", "*.jpg", "*.JPG", "*.PNG"]) + + # --- masks --- + mask_filenames = [] + if self.config.use_masks: + mask_filenames = find_files(f"{split_dir}/mask", exts=["*.png", "*.jpg", "*.JPG", "*.PNG"]) + + dataparser_outputs = DataparserOutputs( + image_filenames=image_filenames, + cameras=cameras, + scene_box=scene_box, + mask_filenames=mask_filenames if len(mask_filenames) > 0 else None, + dataparser_scale=self.config.scale_factor, + ) + return dataparser_outputs diff --git a/nerfstudio/data/dataparsers/phototourism_dataparser.py b/nerfstudio/data/dataparsers/phototourism_dataparser.py index e5a88e10e..d470215ed 100644 --- a/nerfstudio/data/dataparsers/phototourism_dataparser.py +++ b/nerfstudio/data/dataparsers/phototourism_dataparser.py @@ -136,8 +136,8 @@ class Phototourism(DataParser): ) # equally spaced training images starting and ending at 0 and num_images-1 i_eval = np.setdiff1d(i_all, i_train) # eval images are the remaining images i_all = torch.tensor(i_all) - i_train = torch.tensor(i_train) - i_eval = torch.tensor(i_eval) + i_train = torch.tensor(i_train, dtype=torch.long) + i_eval = torch.tensor(i_eval, dtype=torch.long) assert len(i_eval) == num_eval_images if split == "train": indices = i_train diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 97ff9244a..dd34aa7b4 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -138,6 +138,18 @@ class Trainer: 'test': loads train/test datasets into memory 'inference': does not load any dataset into memory """ + # set up viewer if enabled + viewer_log_path = self.base_dir / self.config.viewer.relative_log_filename + self.viewer_state, banner_messages = None, None + if self.config.is_viewer_enabled() and self.local_rank == 0: + datapath = self.config.data + if datapath is None: + datapath = self.base_dir + self.viewer_state, banner_messages = viewer_utils.setup_viewer( + self.config.viewer, log_filename=viewer_log_path, datapath=datapath + ) + self._check_viewer_warnings() + self.pipeline = self.config.pipeline.setup( device=self.device, test_mode=test_mode, @@ -157,17 +169,6 @@ class Trainer: ) ) - # set up viewer if enabled - viewer_log_path = self.base_dir / self.config.viewer.relative_log_filename - self.viewer_state, banner_messages = None, None - if self.config.is_viewer_enabled() and self.local_rank == 0: - datapath = self.pipeline.datamanager.get_datapath() - if datapath is None: - datapath = self.base_dir - self.viewer_state, banner_messages = viewer_utils.setup_viewer( - self.config.viewer, log_filename=viewer_log_path, datapath=datapath - ) - self._check_viewer_warnings() # set up writers/profilers if enabled writer_log_path = self.base_dir / self.config.logging.relative_log_dir writer.setup_event_writer( diff --git a/nerfstudio/field_components/encodings.py b/nerfstudio/field_components/encodings.py index f77c17f4d..462080287 100644 --- a/nerfstudio/field_components/encodings.py +++ b/nerfstudio/field_components/encodings.py @@ -16,8 +16,9 @@ Encoding functions """ +import itertools from abc import abstractmethod -from typing import Optional +from typing import Optional, Sequence import numpy as np import torch @@ -242,7 +243,6 @@ class HashEncoding(Encoding): implementation: Literal["tcnn", "torch"] = "tcnn", interpolation: Optional[Literal["Nearest", "Linear", "Smoothstep"]] = None, ) -> None: - super().__init__(in_dim=3) self.num_levels = num_levels self.features_per_level = features_per_level @@ -557,6 +557,87 @@ class TriplaneEncoding(Encoding): self.resolution = resolution +class KPlanesEncoding(Encoding): + """Learned K-Planes encoding + + A plane encoding supporting both 3D and 4D coordinates. With 3D coordinates this is similar to + :class:`TriplaneEncoding`. With 4D coordinates, the encoding at point ``[i,j,k,q]`` is + a n-dimensional vector computed as the elementwise product of 6 n-dimensional vectors at + ``planes[i,j]``, ``planes[i,k]``, ``planes[i,q]``, ``planes[j,k]``, ``planes[j,q]``, + ``planes[k,q]``. + + Unlike :class:`TriplaneEncoding` this class supports different resolution along each axis. + + This will return a tensor of shape (bs:..., num_components) + + Args: + resolution: Resolution of the grid. Can be a sequence of 3 or 4 integers. + num_components: The number of scalar planes to use (ie: output feature size) + init_a: The lower-bound of the uniform distribution used to initialize the spatial planes + init_b: The upper-bound of the uniform distribution used to initialize the spatial planes + reduce: Whether to use the element-wise product of the planes or the sum + """ + + def __init__( + self, + resolution: Sequence[int] = (128, 128, 128), + num_components: int = 64, + init_a: float = 0.1, + init_b: float = 0.5, + reduce: Literal["sum", "product"] = "product", + ) -> None: + super().__init__(in_dim=len(resolution)) + + self.resolution = resolution + self.num_components = num_components + self.reduce = reduce + if self.in_dim not in {3, 4}: + raise ValueError( + f"The dimension of coordinates must be either 3 (static scenes) " + f"or 4 (dynamic scenes). Found resolution with {self.in_dim} dimensions." + ) + has_time_planes = self.in_dim == 4 + + self.coo_combs = list(itertools.combinations(range(self.in_dim), 2)) + # Unlike the Triplane encoding, we use a parameter list instead of batching all planes + # together to support uneven resolutions (especially useful for time). + # Dynamic models (in_dim == 4) will have 6 planes: + # (y, x), (z, x), (t, x), (z, y), (t, y), (t, z) + # static models (in_dim == 3) will only have the 1st, 2nd and 4th planes. + self.plane_coefs = nn.ParameterList() + for coo_comb in self.coo_combs: + new_plane_coef = nn.Parameter( + torch.empty([self.num_components] + [self.resolution[cc] for cc in coo_comb[::-1]]) + ) + if has_time_planes and 3 in coo_comb: # Time planes initialized to 1 + nn.init.ones_(new_plane_coef) + else: + nn.init.uniform_(new_plane_coef, a=init_a, b=init_b) + self.plane_coefs.append(new_plane_coef) + + def get_out_dim(self) -> int: + return self.num_components + + def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + """Sample features from this encoder. Expects ``in_tensor`` to be in range [-1, 1]""" + original_shape = in_tensor.shape + + output = 1.0 if self.reduce == "product" else 0.0 # identity for corresponding op + for ci, coo_comb in enumerate(self.coo_combs): + grid = self.plane_coefs[ci].unsqueeze(0) # [1, feature_dim, reso1, reso2] + coords = in_tensor[..., coo_comb].view(1, 1, -1, 2) # [1, 1, flattened_bs, 2] + interp = F.grid_sample( + grid, coords, align_corners=True, padding_mode="border" + ) # [1, output_dim, 1, flattened_bs] + interp = interp.view(self.num_components, -1).T # [flattened_bs, output_dim] + if self.reduce == "product": + output = output * interp + else: + output = output + interp + + return output.reshape(*original_shape[:-1], self.num_components) + + class SHEncoding(Encoding): """Spherical harmonic encoding diff --git a/nerfstudio/generative/stable_diffusion_img2img.py b/nerfstudio/generative/stable_diffusion_img2img.py index ec9b96a94..05200f099 100644 --- a/nerfstudio/generative/stable_diffusion_img2img.py +++ b/nerfstudio/generative/stable_diffusion_img2img.py @@ -348,7 +348,7 @@ class StableDiffusionImg2Img(nn.Module): latent = self.imgs_to_latent(image.half()) # t = torch.randint(self.max_step - 1, self.max_step, [1], dtype=torch.long, device=self.device) - t = torch.randint(0, 1, [1], dtype=torch.long) + t = torch.randint(4, 5, [1], dtype=torch.long) new_steps = self.scheduler.timesteps[t:] t_noise = self.scheduler.timesteps[t] @@ -358,7 +358,6 @@ class StableDiffusionImg2Img(nn.Module): # for i, time in enumerate(self.scheduler.timesteps): for i, time in enumerate(new_steps): - # for i in range(num_inference_steps): # predict the noise residual latent_model_input = torch.cat([latent] * 2) @@ -374,7 +373,6 @@ class StableDiffusionImg2Img(nn.Module): latent = self.scheduler.step(noise_pred, time, latent)["prev_sample"] # type: ignore diffused_img = self.latents_to_img(latent.half()) - # diffused_img = F.interpolate(diffused_img, (128, 128), mode="bilinear") diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy() diffused_img = (diffused_img * 255).round().astype("uint8")[0] # might delete this indexing diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index 174c131cc..d2db011d3 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -338,7 +338,7 @@ class MiDaSMSELoss(nn.Module): def __init__(self, reduction_type: Literal["image", "batch"] = "batch"): super().__init__() - self.reduction_type = reduction_type + self.reduction_type: Literal["image", "batch"] = reduction_type # reduction here is different from the image/batch-based reduction. This is either "mean" or "sum" self.mse_loss = MSELoss(reduction="none") diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index d6cb080c6..641da83e1 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -645,6 +645,7 @@ class NeuSSampler(Sampler): # Start with uniform sampling if ray_samples is None: ray_samples = self.uniform_sampler(ray_bundle, num_samples=self.num_samples) + assert ray_samples is not None total_iters = 0 sorted_index = None @@ -701,6 +702,7 @@ class NeuSSampler(Sampler): """ batch_size = ray_samples.shape[0] prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:] + assert ray_samples.deltas is not None deltas = ray_samples.deltas[:, :-1, 0] mid_sdf = (prev_sdf + next_sdf) * 0.5 cos_val = (next_sdf - prev_sdf) / (deltas + 1e-5) @@ -742,6 +744,9 @@ class NeuSSampler(Sampler): ray_samples_2 : ray_samples to merge """ + assert ray_samples_1.spacing_starts is not None and ray_samples_2.spacing_starts is not None + assert ray_samples_1.spacing_ends is not None and ray_samples_2.spacing_ends is not None + assert ray_samples_1.spacing_to_euclidean_fn is not None starts_1 = ray_samples_1.spacing_starts[..., 0] starts_2 = ray_samples_2.spacing_starts[..., 0] diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 3278ff9e5..41c2ec812 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -180,7 +180,7 @@ class SHRenderer(nn.Module): components = components_from_spherical_harmonics(levels=levels, directions=directions) rgb = sh * components[..., None, :] # [..., num_samples, 3, sh_components] - rgb = torch.sum(sh, dim=-1) + 0.5 # [..., num_samples, 3] + rgb = torch.sum(rgb, dim=-1) # [..., num_samples, 3] if self.activation is not None: rgb = self.activation(rgb) diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index 933609884..d6c55856e 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -104,6 +104,7 @@ class Model(nn.Module): # NOTE: call `super().populate_modules()` in subclasses if self.config.enable_collider: + assert self.config.collider_params is not None self.collider = NearFarCollider( near_plane=self.config.collider_params["near_plane"], far_plane=self.config.collider_params["far_plane"], diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index 3b27354d1..0a988b111 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -70,8 +70,10 @@ class TensoRFModelConfig(ModelConfig): """specifies a list of iteration step numbers to perform upsampling""" loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss": 1.0}) """Loss specific weights.""" - num_samples: int = 256 + num_samples: int = 50 """Number of samples in field evaluation""" + num_uniform_samples: int = 200 + """Number of samples in density evaluation""" num_den_components: int = 16 """Number of components in density encoding""" num_color_components: int = 48 @@ -88,6 +90,8 @@ class TensoRFModel(Model): config: TensoRF configuration to instantiate model """ + config: TensoRFModelConfig + def __init__( self, config: TensoRFModelConfig, @@ -116,7 +120,6 @@ class TensoRFModel(Model): def get_training_callbacks( self, training_callback_attributes: TrainingCallbackAttributes ) -> List[TrainingCallback]: - # the callback that we want to run every X iterations after the training iteration def reinitialize_optimizer( self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument @@ -219,8 +222,8 @@ class TensoRFModel(Model): ) # samplers - self.sampler_uniform = UniformSampler(num_samples=self.config.num_samples, single_jitter=True) - self.sampler_pdf = PDFSampler(num_samples=self.config.num_samples // 2, single_jitter=True) + self.sampler_uniform = UniformSampler(num_samples=self.config.num_uniform_samples, single_jitter=True) + self.sampler_pdf = PDFSampler(num_samples=self.config.num_samples, single_jitter=True, include_original=False) # renderers self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) @@ -302,6 +305,7 @@ class TensoRFModel(Model): image = batch["image"].to(outputs["rgb"].device) rgb = outputs["rgb"] acc = colormaps.apply_colormap(outputs["accumulation"]) + assert self.config.collider_params is not None depth = colormaps.apply_depth_colormap( outputs["depth"], accumulation=outputs["accumulation"], diff --git a/nerfstudio/models/vanilla_nerf.py b/nerfstudio/models/vanilla_nerf.py index 53e838891..b69a35358 100644 --- a/nerfstudio/models/vanilla_nerf.py +++ b/nerfstudio/models/vanilla_nerf.py @@ -201,6 +201,7 @@ class NeRFModel(Model): rgb_fine = outputs["rgb_fine"] acc_coarse = colormaps.apply_colormap(outputs["accumulation_coarse"]) acc_fine = colormaps.apply_colormap(outputs["accumulation_fine"]) + assert self.config.collider_params is not None depth_coarse = colormaps.apply_depth_colormap( outputs["depth_coarse"], accumulation=outputs["accumulation_coarse"], diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 62e6c12a5..9c97a88ce 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -97,6 +97,7 @@ class Pipeline(nn.Module): datamanager: DataManager _model: Model generative: bool + world_size: int @property def model(self): @@ -109,9 +110,22 @@ class Pipeline(nn.Module): return self.model.device def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True): - model_state = {key[len("_model.") :]: value for key, value in state_dict.items() if key.startswith("_model.")} + is_ddp_model_state = True + model_state = {} + for key, value in state_dict.items(): + if key.startswith("_model."): + # remove the "_model." prefix from key + model_state[key[len("_model.") :]] = value + # make sure that the "module." prefix comes from DDP, + # rather than an attribute of the model named "module" + if not key.startswith("_model.module."): + is_ddp_model_state = False + # remove "module." prefix added by DDP + if is_ddp_model_state: + model_state = {key[len("module.") :]: value for key, value in model_state.items()} + pipeline_state = {key: value for key, value in state_dict.items() if not key.startswith("_model.")} - self._model.load_state_dict(model_state, strict=strict) + self.model.load_state_dict(model_state, strict=strict) super().load_state_dict(pipeline_state, strict=False) @profiler.time_function @@ -205,6 +219,7 @@ class VanillaPipelineConfig(cfg.InstantiateConfig): class VanillaPipeline(Pipeline): """The pipeline class for the vanilla nerf setup of multiple cameras for one or a few scenes. + Args: config: configuration to instantiate pipeline device: location to place model and data test_mode: @@ -379,7 +394,7 @@ class VanillaPipeline(Pipeline): state = { (key[len("module.") :] if key.startswith("module.") else key): value for key, value in loaded_state.items() } - self._model.update_to_step(step) + self.model.update_to_step(step) self.load_state_dict(state, strict=True) def get_training_callbacks( diff --git a/nerfstudio/pipelines/iterative_pipeline.py b/nerfstudio/pipelines/iterative_pipeline.py index 24d142b6a..fc2d17cd8 100644 --- a/nerfstudio/pipelines/iterative_pipeline.py +++ b/nerfstudio/pipelines/iterative_pipeline.py @@ -101,7 +101,7 @@ class IterativePipeline(VanillaPipeline): @torch.no_grad() def save_new_train_images(self, step: int): print("saving training images") - num_images_per_train = 7 + num_images_per_train = 5 resolution = 512 ray_bundles, cameras, camera_angles = self.datamanager.random_train_views(step, resolution=resolution, num_views=num_images_per_train) train_img_paths = [] @@ -148,10 +148,11 @@ class IterativePipeline(VanillaPipeline): # use a data loader that allows for multiview supervision each train step # rerender train dataset every Y iterations - if step % 2000 == 0: + if step % 500 == 0: self.save_new_train_images(step) if step % 500 == 0: + # if step == 0: self.update_train_images(step) if self.world_size > 1 and step: diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index 859222441..f4b9e9b03 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -55,7 +55,7 @@ def get_colmap_version(colmap_cmd: str, default_version=3.8) -> float: Returns: The version of COLMAP. """ - output = run_command(colmap_cmd, verbose=False) + output = run_command(f"{colmap_cmd} -h", verbose=False) assert output is not None for line in output.split("\n"): if line.startswith("COLMAP"): diff --git a/nerfstudio/process_data/hloc_utils.py b/nerfstudio/process_data/hloc_utils.py index 0aec2376e..f04f6c32c 100644 --- a/nerfstudio/process_data/hloc_utils.py +++ b/nerfstudio/process_data/hloc_utils.py @@ -75,6 +75,11 @@ def run_hloc( camera_model: Camera model to use. gpu: If True, use GPU. verbose: If True, logs the output of the command. + matching_method: Method to use for matching images. + feature_type: Type of visual features to use. + matcher_type: Type of feature matcher to use. + num_matched: Number of image pairs for loc. + refine_pixsfm: If True, refine the reconstruction using pixel-perfect-sfm. """ if not _HAS_HLOC: CONSOLE.print( @@ -112,7 +117,13 @@ def run_hloc( camera_model=camera_model.value ) if refine_pixsfm: - sfm = PixSfM({"dense_features": {"max_edge": 1024}}) + sfm = PixSfM( + conf={ + "dense_features": {"use_cache": True}, + "KA": {"dense_features": {"use_cache": True}, "max_kps_per_problem": 1000}, + "BA": {"strategy": "costmaps"}, + } + ) refined, _ = sfm.reconstruction( sfm_dir, image_dir, diff --git a/nerfstudio/process_data/insta360_utils.py b/nerfstudio/process_data/insta360_utils.py deleted file mode 100644 index 56e554e01..000000000 --- a/nerfstudio/process_data/insta360_utils.py +++ /dev/null @@ -1,190 +0,0 @@ -# Copyright 2022 The Nerfstudio Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Helper utils for processing insta360 data.""" - -import sys -from pathlib import Path -from typing import List, Tuple - -from rich.console import Console - -from nerfstudio.process_data.process_data_utils import get_num_frames_in_video -from nerfstudio.utils.rich_utils import status -from nerfstudio.utils.scripts import run_command - -CONSOLE = Console(width=120) - - -def get_insta360_filenames(data: Path) -> Tuple[Path, Path]: - """Returns the filenames of the Insta360 videos from a single video file. - - Example input name: VID_20220212_070353_00_003.insv - - Args: - data: Path to a Insta360 file. - - Returns: - The filenames of the Insta360 videios. - """ - if data.suffix != ".insv": - raise ValueError("The input file must be an .insv file.") - file_parts = data.stem.split("_") - - stem_back = f"VID_{file_parts[1]}_{file_parts[2]}_00_{file_parts[4]}.insv" - stem_front = f"VID_{file_parts[1]}_{file_parts[2]}_10_{file_parts[4]}.insv" - - filename_back = data.parent / stem_back - filename_front = data.parent / stem_front - - return filename_back, filename_front - - -def convert_insta360_to_images( - video_front: Path, - video_back: Path, - image_dir: Path, - num_frames_target: int, - crop_percentage: float = 0.7, - verbose: bool = False, -) -> Tuple[List[str], int]: - """Converts a video into a sequence of images. - - Args: - video_front: Path to the front video. - video_back: Path to the back video. - output_dir: Path to the output directory. - num_frames_target: Number of frames to extract. - crop_percentage: Percentage used to calculate the cropped dimensions of extracted frames. Currently used to crop - out the curved portions of the fish-eye lens. - verbose: If True, logs the output of the command. - Returns: - A tuple containing summary of the conversion and the number of extracted frames. - """ - - with status(msg="Converting video to images...", spinner="bouncingBall", verbose=verbose): - # delete existing images in folder - for img in image_dir.glob("*.png"): - if verbose: - CONSOLE.log(f"Deleting {img}") - img.unlink() - - num_frames_front = get_num_frames_in_video(video_front) - num_frames_back = get_num_frames_in_video(video_back) - if num_frames_front == 0: - CONSOLE.print(f"[bold red]Error: Video has no frames: {video_front}") - sys.exit(1) - if num_frames_back == 0: - CONSOLE.print(f"[bold red]Error: Video has no frames: {video_front}") - sys.exit(1) - - spacing = num_frames_front // (num_frames_target // 2) - vf_cmds = [] - if spacing > 1: - vf_cmds = [f"thumbnail={spacing}", "setpts=N/TB"] - else: - CONSOLE.print("[bold red]Can't satisfy requested number of frames. Extracting all frames.") - - vf_cmds.append(f"crop=iw*{crop_percentage}:ih*{crop_percentage}") - - front_vf_cmds = vf_cmds + ["transpose=2"] - back_vf_cmds = vf_cmds + ["transpose=1"] - - front_ffmpeg_cmd = f"ffmpeg -i {video_front} -vf {','.join(front_vf_cmds)} -r 1 {image_dir / 'frame_%05d.png'}" - back_ffmpeg_cmd = ( - f"ffmpeg -i {video_back} -vf {','.join(back_vf_cmds)} -r 1 {image_dir / 'back_frame_%05d.png'}" - ) - - run_command(front_ffmpeg_cmd, verbose=verbose) - run_command(back_ffmpeg_cmd, verbose=verbose) - - num_extracted_front_frames = len(list(image_dir.glob("frame*.png"))) - for i, img in enumerate(image_dir.glob("back_frame_*.png")): - img.rename(image_dir / f"frame_{i+1+num_extracted_front_frames:05d}.png") - - num_final_frames = len(list(image_dir.glob("*.png"))) - summary_log = [] - summary_log.append(f"Starting with {num_frames_front + num_frames_back} video frames") - summary_log.append(f"We extracted {num_final_frames} images") - CONSOLE.log("[bold green]:tada: Done converting insta360 to images.") - - return summary_log, num_final_frames - - -def convert_insta360_single_file_to_images( - video: Path, - image_dir: Path, - num_frames_target: int, - crop_percentage: float = 0.7, - verbose: bool = False, -) -> Tuple[List[str], int]: - """Converts a video into a sequence of images. - - Args: - video: Path to the video. - output_dir: Path to the output directory. - num_frames_target: Number of frames to extract. - crop_percentage: Percentage used to calculate the cropped dimensions of extracted frames. Currently used to crop - out the curved portions of the fish-eye lens. - verbose: If True, logs the output of the command. - Returns: - A tuple containing summary of the conversion and the number of extracted frames. - """ - - with status(msg="Converting video to images...", spinner="bouncingBall", verbose=verbose): - # delete existing images in folder - for img in image_dir.glob("*.png"): - if verbose: - CONSOLE.log(f"Deleting {img}") - img.unlink() - - num_frames = get_num_frames_in_video(video) - if num_frames == 0: - CONSOLE.print(f"[bold red]Error: Video has no frames: {video}") - sys.exit(1) - - spacing = num_frames // (num_frames_target // 2) - vf_cmds = [] - if spacing > 1: - vf_cmds = [f"thumbnail={spacing}", "setpts=N/TB"] - else: - CONSOLE.print("[bold red]Can't satisfy requested number of frames. Extracting all frames.") - - vf_cmds_back = vf_cmds.copy() - vf_cmds_front = vf_cmds.copy() - - vf_cmds_back.append( - f"crop=ih*{crop_percentage}:ih*{crop_percentage}:ih*{crop_percentage}/4:ih*{crop_percentage}/4" - ) - vf_cmds_front.append( - f"crop=ih*{crop_percentage}:ih*{crop_percentage}:iw/2+ih*{crop_percentage/4}:ih*{crop_percentage/4}" - ) - - front_ffmpeg_cmd = f"ffmpeg -i {video} -vf {','.join(vf_cmds_front)} -r 1 {image_dir / 'frame_%05d.png'}" - back_ffmpeg_cmd = f"ffmpeg -i {video} -vf {','.join(vf_cmds_back)} -r 1 {image_dir / 'back_frame_%05d.png'}" - - run_command(back_ffmpeg_cmd, verbose=verbose) - run_command(front_ffmpeg_cmd, verbose=verbose) - - num_extracted_frames = len(list(image_dir.glob("frame*.png"))) - for i, img in enumerate(image_dir.glob("back_frame_*.png")): - img.rename(image_dir / f"frame_{i+1+num_extracted_frames:05d}.png") - - num_final_frames = len(list(image_dir.glob("*.png"))) - summary_log = [] - summary_log.append(f"Starting with {num_frames} video frames") - summary_log.append(f"We extracted {num_final_frames} images") - CONSOLE.log("[bold green]:tada: Done converting insta360 to images.") - - return summary_log, num_final_frames diff --git a/nerfstudio/process_data/metashape_utils.py b/nerfstudio/process_data/metashape_utils.py index 17d7771e8..0e4a7593e 100644 --- a/nerfstudio/process_data/metashape_utils.py +++ b/nerfstudio/process_data/metashape_utils.py @@ -27,7 +27,7 @@ from nerfstudio.process_data.process_data_utils import CAMERA_MODELS CONSOLE = Console(width=120) -def _find_distortion_param(calib_xml: ET.Element, param_name: str): +def _find_param(calib_xml: ET.Element, param_name: str): param = calib_xml.find(param_name) if param is not None: return float(param.text) # type: ignore @@ -62,34 +62,71 @@ def metashape_to_json( # pylint: disable=too-many-statements raise ValueError("No sensors found") calibrated_sensors = [sensor for sensor in sensors if sensor.find("calibration")] - if len(calibrated_sensors) != 1: - raise ValueError("Only one sensor is supported for now") - - sensor = calibrated_sensors[0] - + if not calibrated_sensors: + raise ValueError("No calibrated sensor found in Metashape XML") + sensor_type = [s.get("type") for s in calibrated_sensors] + if sensor_type.count(sensor_type[0]) != len(sensor_type): + raise ValueError("All Metashape sensors do not have the same sensor type") data = {} + if sensor_type[0] == "frame": + data["camera_model"] = CAMERA_MODELS["perspective"].value + elif sensor_type[0] == "fisheye": + data["camera_model"] = CAMERA_MODELS["fisheye"].value + elif sensor_type[0] == "spherical": + data["camera_model"] = CAMERA_MODELS["equirectangular"].value + else: + # Cylindrical and RPC sensor types are not supported + raise ValueError(f"Unsupported Metashape sensor type '{sensor_type[0]}'") - assert sensor is not None, "Sensor not found in Metashape XML" - resolution = sensor.find("resolution") - assert resolution is not None, "Resolution not found in Metashape xml" - data["w"] = int(resolution.get("width")) # type: ignore - data["h"] = int(resolution.get("height")) # type: ignore + sensor_dict = {} + for sensor in calibrated_sensors: + s = {} + resolution = sensor.find("resolution") + assert resolution is not None, "Resolution not found in Metashape xml" + s["w"] = int(resolution.get("width")) # type: ignore + s["h"] = int(resolution.get("height")) # type: ignore - calib = sensor.find("calibration") - assert calib is not None, "Calibration not found in Metashape xml" - data["fl_x"] = float(calib.find("f").text) # type: ignore - data["fl_y"] = float(calib.find("f").text) # type: ignore - data["cx"] = float(calib.find("cx").text) + data["w"] / 2.0 # type: ignore - data["cy"] = float(calib.find("cy").text) + data["h"] / 2.0 # type: ignore + calib = sensor.find("calibration") + f = calib.find("f") + if f is not None: + s["fl_x"] = s["fl_y"] = float(f.text) # type: ignore + s["cx"] = _find_param(calib, "cx") + s["w"] / 2.0 # type: ignore + s["cy"] = _find_param(calib, "cy") + s["h"] / 2.0 # type: ignore - data["k1"] = _find_distortion_param(calib, "k1") - data["k2"] = _find_distortion_param(calib, "k2") - data["k3"] = _find_distortion_param(calib, "k3") - data["k4"] = _find_distortion_param(calib, "k4") - data["p1"] = _find_distortion_param(calib, "p1") - data["p2"] = _find_distortion_param(calib, "p2") + s["k1"] = _find_param(calib, "k1") + s["k2"] = _find_param(calib, "k2") + s["k3"] = _find_param(calib, "k3") + s["k4"] = _find_param(calib, "k4") + s["p1"] = _find_param(calib, "p1") + s["p2"] = _find_param(calib, "p2") - data["camera_model"] = CAMERA_MODELS["perspective"].value + sensor_dict[sensor.get("id")] = s + + components = chunk.find("components") + component_dict = {} + for component in components: + transform = component.find("transform") + if transform is not None: + rotation = transform.find("rotation") + if rotation is None: + r = np.eye(3) + else: + r = np.array([float(x) for x in rotation.text.split()]).reshape((3, 3)) + translation = transform.find("translation") + if translation is None: + t = np.zeros(3) + else: + t = np.array([float(x) for x in translation.text.split()]) + scale = transform.find("scale") + if scale is None: + s = 1.0 + else: + s = float(scale.text) + + m = np.eye(4) + m[:3, :3] = r + m[:3, 3] = t / s + component_dict[component.get("id")] = m frames = [] cameras = chunk.find("cameras") @@ -97,33 +134,36 @@ def metashape_to_json( # pylint: disable=too-many-statements num_skipped = 0 for camera in cameras: frame = {} - # Labels sometimes have a file extension. We remove it for consistency. - camera_label = camera.get("label").split(".")[0] # type: ignore + camera_label = camera.get("label") if camera_label not in image_filename_map: - continue + # Labels sometimes have a file extension. Try without the extension. + # (maybe it's just a '.' in the image name) + camera_label = camera_label.split(".")[0] # type: ignore + if camera_label not in image_filename_map: + continue frame["file_path"] = image_filename_map[camera_label].as_posix() - if camera.get("sensor_id") != sensor.get("id"): + sensor_id = camera.get("sensor_id") + if sensor_id not in sensor_dict: # this should only happen when we have a sensor that doesn't have calibration if verbose: CONSOLE.print(f"Missing sensor calibration for {camera.get('label')}, Skipping") num_skipped += 1 continue + # Add all sensor parameters to this frame. + frame.update(sensor_dict[sensor_id]) if camera.find("transform") is None: if verbose: CONSOLE.print(f"Missing transforms data for {camera.get('label')}, Skipping") num_skipped += 1 continue - t = [float(x) for x in camera.find("transform").text.split()] # type: ignore - transform = np.array( - [ - [t[8], -t[9], -t[10], t[11]], - [t[0], -t[1], -t[2], t[3]], - [t[4], -t[5], -t[6], t[7]], - [t[12], -t[13], -t[14], t[15]], - ] - ) + transform = np.array([float(x) for x in camera.find("transform").text.split()]).reshape((4, 4)) + component_id = camera.get("component_id") + if component_id in component_dict: + transform = component_dict[component_id] @ transform + transform = transform[[2, 0, 1, 3], :] + transform[:, 1:3] *= -1 frame["transform_matrix"] = transform.tolist() frames.append(frame) diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index a73127fcb..a977df8b2 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -14,6 +14,7 @@ """Helper utils for processing data into the nerfstudio format.""" +import math import os import shutil import sys @@ -141,7 +142,7 @@ def convert_video_to_images( if num_frames == 0: CONSOLE.print(f"[bold red]Error: Video has no frames: {video_path}") sys.exit(1) - print("Number of frames in video:", num_frames) + CONSOLE.print("Number of frames in video:", num_frames) out_filename = image_dir / "frame_%05d.png" ffmpeg_cmd = f'ffmpeg -i "{video_path}"' @@ -157,9 +158,12 @@ def convert_video_to_images( spacing = num_frames // num_frames_target if spacing > 1: ffmpeg_cmd += f" -vf thumbnail={spacing},setpts=N/TB{crop_cmd} -r 1" + CONSOLE.print("Number of frames to extract:", math.ceil(num_frames / spacing)) else: CONSOLE.print("[bold red]Can't satisfy requested number of frames. Extracting all frames.") - ffmpeg_cmd += f" -pix_fmt bgr8 -vf {crop_cmd[1:]}" + ffmpeg_cmd += " -pix_fmt bgr8" + if crop_cmd != "": + ffmpeg_cmd += f" -vf {crop_cmd[1:]}" ffmpeg_cmd += f" {out_filename}" run_command(ffmpeg_cmd, verbose=verbose) diff --git a/nerfstudio/process_data/realitycapture_utils.py b/nerfstudio/process_data/realitycapture_utils.py index 7bc73e132..918aa84a5 100644 --- a/nerfstudio/process_data/realitycapture_utils.py +++ b/nerfstudio/process_data/realitycapture_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Helper utils for processing polycam data into the nerfstudio format.""" +"""Helper utils for processing reality capture data into the nerfstudio format.""" import csv import json @@ -59,27 +59,21 @@ def realitycapture_to_json( for column, value in row.items(): cameras.setdefault(column, []).append(value) - for name in cameras["#name"]: - camera_label = name.split(".")[0] - if camera_label in image_filename_map: - img = np.array(Image.open(output_dir / image_filename_map[camera_label])) - break - - height, width, _ = img.shape - - data["h"] = int(height) - data["w"] = int(width) - missing_image_data = 0 for i, name in enumerate(cameras["#name"]): - basename = name.split(".")[0] + basename = name.rpartition(".")[0] if basename not in image_filename_map: if verbose: CONSOLE.print(f"Missing image for camera data {basename}, Skipping") missing_image_data += 1 continue + frame = {} + img = np.array(Image.open(output_dir / image_filename_map[basename])) + height, width, _ = img.shape + frame["h"] = int(height) + frame["w"] = int(width) frame["file_path"] = image_filename_map[basename].as_posix() frame["fl_x"] = float(cameras["f"][i]) * max(width, height) / 36 frame["fl_y"] = float(cameras["f"][i]) * max(width, height) / 36 diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index 83d06e371..bffb0bd03 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -283,7 +283,7 @@ class WandbWriter(Writer): """WandDB Writer Class""" def __init__(self, log_dir: Path): - wandb.init(project="dream-embed", dir=str(log_dir), reinit=True) + wandb.init(project="nerfstudio-project", dir=str(log_dir), reinit=True, name=log_dir.name) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: image = torch.permute(image, (2, 0, 1)) diff --git a/nerfstudio/viewer/app/src/index.jsx b/nerfstudio/viewer/app/src/index.jsx index 966e50e43..2d16f4a01 100644 --- a/nerfstudio/viewer/app/src/index.jsx +++ b/nerfstudio/viewer/app/src/index.jsx @@ -1,12 +1,12 @@ import './index.scss'; import React from 'react'; -import ReactDOM from 'react-dom'; +import { createRoot } from 'react-dom/client'; import { Provider } from 'react-redux'; import App from './App'; import WebSocketProvider from './modules/WebSocket/WebSocket'; import store from './store'; -const root = ReactDOM.createRoot(document.getElementById('root')); +const root = createRoot(document.getElementById('root')); root.render( diff --git a/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx b/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx index 2bb6d8fd3..f378c5614 100644 --- a/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx @@ -1,6 +1,8 @@ -import { buttonGroup, useControls } from 'leva'; +import * as React from 'react'; +import { Leva, buttonGroup, useControls } from 'leva'; import { useContext, useEffect, useState } from 'react'; import { useDispatch, useSelector } from 'react-redux'; +import LevaTheme from '../../themes/leva_theme.json'; import { WebSocketContext } from '../WebSocket/WebSocket'; @@ -22,7 +24,7 @@ function dispatch_and_send(websocket, dispatch, path, data) { } } -export function RenderControls() { +function ControlsLeva() { // connection status indicators const websocket = useContext(WebSocketContext).socket; const outputOptions = useSelector( @@ -360,3 +362,18 @@ export function RenderControls() { return null; } + +export function RenderControls() { + return ( +
+ + +
+ ); +} diff --git a/nerfstudio/viewer/app/src/modules/LoadPathModal/LoadPathModal.jsx b/nerfstudio/viewer/app/src/modules/LoadPathModal/LoadPathModal.jsx index 5f79d1e00..5b5d2fced 100644 --- a/nerfstudio/viewer/app/src/modules/LoadPathModal/LoadPathModal.jsx +++ b/nerfstudio/viewer/app/src/modules/LoadPathModal/LoadPathModal.jsx @@ -113,9 +113,9 @@ export default function LoadPathModal(props: LoadPathModalProps) {
-

-

OR
-

+
+

OR

+
)} {camera_paths_arr.length === 0 && ( diff --git a/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx b/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx index 8388f7735..ef5a38888 100644 --- a/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx +++ b/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx @@ -31,10 +31,10 @@ export function get_scene_tree() { }, setValue(key, value) { this.value.set(key, value); - this.callbacks.forEach((cbk) => { - const i = cbk[1]; - const callback = cbk[0]; - return callback(this.value.get(i)); + this.callbacks.forEach((callback, callbackKey) => { + if (callbackKey === key) { + callback(value); + } }); }, }; @@ -164,6 +164,9 @@ export function get_scene_tree() { if (!scene_state.value.get('mouse_in_scene')) { return; } + if (keyMap.Space === true) { + camera_controls.setLookAt(0.7, -0.7, 0.3, 0, 0, 0); + } translate(); rotate(); } @@ -189,12 +192,6 @@ export function get_scene_tree() { window.addEventListener('keydown', onKeyDown, true); window.addEventListener('keyup', onKeyUp, true); - window.addEventListener('keydown', (event) => { - if (event.code === 'Space') { - camera_controls.setLookAt(0.7, -0.7, 0.3, 0, 0, 0); - } - }); - sceneTree.metadata.camera_controls = camera_controls; sceneTree.metadata.moveCamera = moveCamera; @@ -390,7 +387,7 @@ export function SceneTreeWebSocketListener() { const dispatch = useDispatch(); useEffect(() => { - socket.addEventListener('message', (originalCmd) => { + const messageHandler = (originalCmd) => { // set the remote description when the offer is received const cmd = msgpack.decode(new Uint8Array(originalCmd.data)); if (cmd.type === 'write') { @@ -401,6 +398,11 @@ export function SceneTreeWebSocketListener() { data: cmd.data, }); } - }); + }; + socket.addEventListener('message', messageHandler); + + return () => { + socket.removeEventListener('message', messageHandler); + }; }, [socket]); // dependency to call this whenever the websocket changes } diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx index 478068e6f..dce2aad0b 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx @@ -147,9 +147,7 @@ function RenderTimeSelector(props) { disabled={disabled} sx={{ input: { - '-webkit-text-fill-color': `${ - disabled ? '#24B6FF' : '#EBEBEB' - } !important`, + WebkitTextFillColor: `${disabled ? '#24B6FF' : '#EBEBEB'} !important`, color: `${disabled ? '#24B6FF' : '#EBEBEB'} !important`, }, }} @@ -282,9 +280,7 @@ function FovSelector(props) { }} sx={{ input: { - '-webkit-text-fill-color': `${ - disabled ? '#24B6FF' : '#EBEBEB' - } !important`, + WebkitTextFillColor: `${disabled ? '#24B6FF' : '#EBEBEB'} !important`, color: `${disabled ? '#24B6FF' : '#EBEBEB'} !important`, }, }} @@ -655,13 +651,11 @@ export default function CameraPanel(props) { }; const setCropEnabled = (value) => { - console.log('setting the crop enabled value to: ', value); dispatch({ type: 'write', path: 'renderingState/crop_enabled', data: value, }); - console.log('crop enabled value is now: ', crop_enabled); }; const serCropBgColor = (value) => { @@ -1121,7 +1115,6 @@ export default function CameraPanel(props) { reset_slider_render_on_add(new_camera_list); if ('crop' in camera_path_object && camera_path_object.crop !== null) { - console.log('crop' in camera_path_object); setCropEnabled(true); serCropBgColor(camera_path_object.crop.crop_bg_color); setCropCenter(camera_path_object.crop.crop_center); diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/ExportPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/ExportPanel.jsx index ac998daa2..30e390b34 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/ExportPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/ExportPanel.jsx @@ -30,7 +30,7 @@ function TabPanel(props: TabPanelProps) { > {value === index && ( - {children} + {children} )} @@ -49,9 +49,9 @@ const CLIPPING_BOX_NAME = 'Clipping Box'; export default function ExportPanel(props) { // unpack relevant information const sceneTree = props.sceneTree; - const showExportBox = props.showExportBox; + const showExportBox = useSelector((state) => state.show_export_box); - const [type_value, setTypeValue] = React.useState(0); + const [type_value, setTypeValue] = React.useState(false); const handleTypeChange = (event: React.SyntheticEvent, newValue: number) => { setTypeValue(newValue); diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/MeshSubPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/MeshSubPanel.jsx index 3eb2725af..f345652e8 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/MeshSubPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/MeshSubPanel.jsx @@ -132,9 +132,6 @@ export default function MeshSubPanel(props) { return boolean ? 'True' : 'False'; }; - console.log(controlValues.outputDir); - console.log(getPythonBool(controlValues.removeOutliers)); - const bbox_min = [ clippingCenter[0] - clippingScale[0] / 2, clippingCenter[1] - clippingScale[1] / 2, @@ -147,7 +144,6 @@ export default function MeshSubPanel(props) { ]; const mesh_method_choice = controlValues.mesh_method_options; - console.log(mesh_method_choice); let cmd = ''; if (mesh_method_choice === 'tsdf') { cmd = @@ -192,7 +188,7 @@ export default function MeshSubPanel(props) { inner_html = (
You need to train a model with normals to use poisson meshing. Try tsdf - or retrain the model with normals. You can use nerfacto with
+ or retrain the model with normals. You can use nerfacto with
--pipeline.model.predict-normals True
); diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/SidePanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/SidePanel.jsx index 35dd126d6..292c88a85 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/SidePanel.jsx +++ b/nerfstudio/viewer/app/src/modules/SidePanel/SidePanel.jsx @@ -5,19 +5,16 @@ import * as THREE from 'three'; import Box from '@mui/material/Box'; import Divider from '@mui/material/Divider'; -import { Leva } from 'leva'; import Tab from '@mui/material/Tab'; import Tabs from '@mui/material/Tabs'; -import Typography from '@mui/material/Typography'; import { CameraAltRounded, TuneRounded, WidgetsRounded, ImportExportRounded, } from '@mui/icons-material/'; -import { useSelector } from 'react-redux'; +import { useDispatch, useSelector } from 'react-redux'; import StatusPanel from './StatusPanel'; -import LevaTheme from '../../themes/leva_theme.json'; import CameraPanel from './CameraPanel'; import ScenePanel from './ScenePanel'; import { RenderControls } from '../ConfigPanel/ConfigPanel'; @@ -39,31 +36,6 @@ export const snap_to_camera = (sceneTree, camera, matrix) => { ); }; -interface TabPanelProps { - children: React.ReactNode; - index: number; - value: number; -} - -function TabPanel(props: TabPanelProps) { - const { children, value, index, ...other } = props; - - return ( - - ); -} - function a11yProps(index: number) { return { id: `simple-tab-${index}`, @@ -71,6 +43,86 @@ function a11yProps(index: number) { }; } +interface PanelTabContentsProps { + children: React.ReactNode; + index: number; + value: number; +} + +/** One tab in the control panel. */ +function PanelTabContents(props: PanelTabContentsProps) { + const { children, value, index, ...other } = props; + + return ( + + ); +} + +interface PanelContentsProps { + children: React.ReactNode; +} + +function PanelContents(props: PanelContentsProps) { + const dispatch = useDispatch(); + const [tabState, setTabState] = React.useState(0); + const handleChange = (_event: React.SyntheticEvent, newValue: number) => { + setTabState(newValue); + dispatch({ + type: 'write', + path: 'show_export_box', + data: newValue === 3, + }); + }; + const camera_choice = useSelector( + (state) => state.renderingState.camera_choice, + ); + const arrayChildren = React.Children.toArray(props.children); + + return ( + <> + + + } label="Controls" {...a11yProps(0)} /> + } label="Render" {...a11yProps(1)} /> + } + label="Scene" + disabled={camera_choice === 'Render Camera'} + {...a11yProps(2)} + /> + } + label="Export" + disabled={camera_choice === 'Render Camera'} + {...a11yProps(3)} + /> + + + + {arrayChildren.map((child, index) => ( + // eslint-disable-next-line react/no-array-index-key + + {child} + + ))} + + ); +} + interface BasicTabsProps { sceneTree: object; } @@ -78,78 +130,17 @@ interface BasicTabsProps { export function BasicTabs(props: BasicTabsProps) { const sceneTree = props.sceneTree; - const [value, setValue] = React.useState(0); - const [showExportBox, setShowExportBox] = React.useState(false); - - const handleChange = (event: React.SyntheticEvent, newValue: number) => { - setValue(newValue); - setShowExportBox(newValue === 3); - }; - const camera_choice = useSelector( - (state) => state.renderingState.camera_choice, - ); - - React.useEffect(() => { - if (camera_choice === 'Render Camera') { - setValue(1); - } - }, [camera_choice]); - return (
- - - } label="Controls" {...a11yProps(0)} /> - } label="Render" {...a11yProps(1)} /> - } - label="Scene" - disabled={camera_choice === 'Render Camera'} - {...a11yProps(2)} - /> - } - label="Export" - disabled={camera_choice === 'Render Camera'} - {...a11yProps(3)} - /> - - - -
- - -
- - - - - -
- -
- - - - - + + + + + +
); diff --git a/nerfstudio/viewer/app/src/modules/ViewerWindow/ViewerWindow.jsx b/nerfstudio/viewer/app/src/modules/ViewerWindow/ViewerWindow.jsx index ea6196a14..7f38ad584 100644 --- a/nerfstudio/viewer/app/src/modules/ViewerWindow/ViewerWindow.jsx +++ b/nerfstudio/viewer/app/src/modules/ViewerWindow/ViewerWindow.jsx @@ -162,24 +162,22 @@ export default function ViewerWindow(props) { }; useEffect(() => { - function update_dimensions() { + const handleNewDimensions = () => { setDimensions({ height: get_window_height(), width: get_window_width(), }); - } + }; - window.addEventListener('resize', update_dimensions); - - // force a rerender setDimensions({ height: get_window_height(), width: get_window_width(), }); render(); + window.addEventListener('resize', handleNewDimensions); return () => { - window.removeEventListener('resize', update_dimensions); + window.removeEventListener('resize', handleNewDimensions); }; }, []); diff --git a/nerfstudio/viewer/app/src/reducer.js b/nerfstudio/viewer/app/src/reducer.js index 5aefe9c8a..21279a547 100644 --- a/nerfstudio/viewer/app/src/reducer.js +++ b/nerfstudio/viewer/app/src/reducer.js @@ -13,6 +13,8 @@ const initialState = { render_img: null, // The rendered images + show_export_box: false, // whether to show the export box + // the rendering state renderingState: { // cameras diff --git a/nerfstudio/viewer/app/src/subscriber.js b/nerfstudio/viewer/app/src/subscriber.js index 5af1a946c..dd95b196f 100644 --- a/nerfstudio/viewer/app/src/subscriber.js +++ b/nerfstudio/viewer/app/src/subscriber.js @@ -1,19 +1,17 @@ import { useContext } from 'react'; import { ReactReduxContext } from 'react-redux'; -export function subscribe_to_changes(selector_fn, fn) { - // selector_fn: returns a value from the redux state - // fn_valid: function to run on a valid input - // fn_null: function to run on a null input +export function subscribe_to_changes(selector_fn, callback_fn) { const { store } = useContext(ReactReduxContext); - let current; + let previousState; const handleChange = () => { - const previous = current; - current = selector_fn(store.getState()); - if (previous !== current) { - fn(previous, current); + const currentState = selector_fn(store.getState()); + if (previousState !== currentState) { + callback_fn(previousState, currentState); + previousState = currentState; } }; + store.subscribe(handleChange); } diff --git a/nerfstudio/viewer/server/subprocess.py b/nerfstudio/viewer/server/subprocess.py index 3e07e2d6b..17a4cdb7a 100644 --- a/nerfstudio/viewer/server/subprocess.py +++ b/nerfstudio/viewer/server/subprocess.py @@ -33,6 +33,43 @@ from nerfstudio.viewer.server import server CONSOLE = Console() +def is_port_open(port: int): + """Returns True if the port is open. + + Args: + port: Port to check. + + Returns: + True if the port is open, False otherwise. + """ + try: + sock = socket.socket() + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + _ = sock.bind(("", port)) + sock.close() + return True + except OSError: + return False + + +def get_free_port(default_port: int = None): + """Returns a free port on the local machine. Try to use default_port is possible. + + Args: + default_port: Port to try to use. + + Returns: + A free port on the local machine. + """ + if default_port: + if is_port_open(default_port): + return default_port + sock = socket.socket() + sock.bind(("", 0)) + port = sock.getsockname()[1] + return port + + def run_viewer_bridge_server_as_subprocess( websocket_port: int, zmq_port: Optional[int] = None, @@ -54,9 +91,7 @@ def run_viewer_bridge_server_as_subprocess( # find an available port for zmq if zmq_port is None: - sock = socket.socket() - sock.bind(("", 0)) - zmq_port = sock.getsockname()[1] + zmq_port = get_free_port() string = f"Using ZMQ port: {zmq_port}" CONSOLE.print(f"[bold yellow]{string}") @@ -90,7 +125,6 @@ def run_viewer_bridge_server_as_subprocess( "You likely have to modify --viewer.zmq-port and/or --viewer.websocket-port in the " "config to avoid conflicting ports.\n" ) - string += "Try modifying --viewer.websocket-port 7007\n" CONSOLE.print(f"[bold red]{string}") cleanup(process) # This exists the entire program. sys.exit() will only kill the thread that this runs in. diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 89d66766f..66c2c3dbe 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -45,7 +45,10 @@ from nerfstudio.utils import colormaps, profiler, writer from nerfstudio.utils.decorators import check_main_thread, decorate_all from nerfstudio.utils.io import load_from_json, write_to_json from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter -from nerfstudio.viewer.server.subprocess import run_viewer_bridge_server_as_subprocess +from nerfstudio.viewer.server.subprocess import ( + get_free_port, + run_viewer_bridge_server_as_subprocess, +) from nerfstudio.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h from nerfstudio.viewer.server.visualizer import Viewer @@ -256,10 +259,13 @@ class ViewerState: self.datapath = datapath.parent if datapath.is_file() else datapath if self.config.launch_bridge_server: # start the viewer bridge server - assert self.config.websocket_port is not None + if self.config.websocket_port is None: + websocket_port = get_free_port(default_port=self.config.websocket_port_default) + else: + websocket_port = self.config.websocket_port self.log_filename.parent.mkdir(exist_ok=True) zmq_port = run_viewer_bridge_server_as_subprocess( - self.config.websocket_port, + websocket_port, zmq_port=self.config.zmq_port, ip_address=self.config.ip_address, log_filename=str(self.log_filename), @@ -267,7 +273,7 @@ class ViewerState: # TODO(ethan): log the output of the viewer bridge server in a file where the training logs go CONSOLE.line() version = get_viewer_version() - websocket_url = f"ws://localhost:{self.config.websocket_port}" + websocket_url = f"ws://localhost:{websocket_port}" self.viewer_url = f"https://viewer.nerf.studio/versions/{version}/?websocket_url={websocket_url}" CONSOLE.rule(characters="=") CONSOLE.print(f"[Public] Open the viewer at {self.viewer_url}") @@ -358,15 +364,6 @@ class ViewerState: max_scene_box = torch.max(dataset.scene_box.aabb[1] - dataset.scene_box.aabb[0]).item() self.vis["renderingState/max_box_size"].write(max_scene_box) - # self.vis["renderingState/render_time"].write(str(0)) - - # set the properties of the camera - # self.vis["renderingState/camera"].write(json_) - - # set the main camera intrinsics to one from the dataset - # K = camera.get_intrinsics_matrix() - # set_persp_intrinsics_matrix(self.vis, K.double().numpy()) - def _check_camera_path_payload(self, trainer, step: int): """Check to see if the camera path export button was pressed.""" # check if we should interrupt from a button press? @@ -679,8 +676,17 @@ class ViewerState: image = selected_output[..., [2, 1, 0]].cpu().numpy() - data = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, 75])[1].tobytes() - data = str("data:image/jpeg;base64," + base64.b64encode(data).decode("ascii")) + data = cv2.imencode( + f".{self.config.image_format}", + image, + [ + cv2.IMWRITE_JPEG_QUALITY, + self.config.jpeg_quality, + cv2.IMWRITE_PNG_COMPRESSION, + self.config.png_compression, + ], + )[1].tobytes() + data = str(f"data:image/{self.config.image_format};base64," + base64.b64encode(data).decode("ascii")) self.vis["render_img"].write(data) def _update_viewer_stats(self, render_time: float, num_rays: int, image_height: int, image_width: int) -> None: diff --git a/pyproject.toml b/pyproject.toml index e8dbf3e96..c22bafe7d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,7 @@ dependencies = [ "rich>=12.5.1", "scikit-image>=0.19.3", "tensorboard==2.9.0", - "torch>=1.12.1", + "torch>=1.12.1,<2.0.0", "torchmetrics[image]>=0.9.3", "torchtyping>=0.1.4", "torchvision>=0.13.0", @@ -76,7 +76,7 @@ dev = [ "pylint==2.13.4", "pytest==7.1.2", "pytest-xdist==2.5.0", - "typeguard>=2.13.3", + "typeguard==2.13.3", ] # Documentation related packages diff --git a/scripts/completions/install.py b/scripts/completions/install.py index 98e67fea2..067bea2bd 100644 --- a/scripts/completions/install.py +++ b/scripts/completions/install.py @@ -399,3 +399,6 @@ def entrypoint(): if __name__ == "__main__": entrypoint() + +# For sphinx docs +get_parser_fn = lambda: tyro.extras.get_parser(main) # noqa diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index 3339b50bc..a92798bbb 100644 --- a/scripts/downloads/download_data.py +++ b/scripts/downloads/download_data.py @@ -359,6 +359,65 @@ class SDFstudioDemoDownload(DatasetDownload): os.remove(download_path) +# pylint: disable=line-too-long +nerfosr_downloads = { + "europa": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=europa&downloadStartSecret=0k2r95c1fdej", + "lk2": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=lk2&downloadStartSecret=w8kuvjzmchc", + "lwp": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=lwp&downloadStartSecret=gtnc4vmkcjq", + "rathaus": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=rathaus&downloadStartSecret=7372aewy6rr", + "schloss": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=schloss&downloadStartSecret=y8t00nqx0h", + "st": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=st&downloadStartSecret=kl9ptuxe8v", + "stjacob": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=stjacob&downloadStartSecret=sntsim6ebvm", + "stjohann": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=stjohann&downloadStartSecret=g80ug1fsbmh", + "trevi": "https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk/download?path=%2FData&files=trevi&downloadStartSecret=ot1483bigjm", + "all": None, +} + + +NeRFOSRCaptureName = tyro.extras.literal_type_from_choices(nerfosr_downloads.keys()) + + +@dataclass +class NeRFOSRDownload(DatasetDownload): + """Download the NeRF-OSR dataset.""" + + capture_name: NeRFOSRCaptureName = "europa" + + def download(self, save_dir: Path): + """Download the NeRF-OSR dataset: https://nextcloud.mpi-klsb.mpg.de/index.php/s/mGXYKpD8raQ8nMk""" + + if self.capture_name == "all": + for capture_name in nerfosr_downloads: + if capture_name != "all": + NeRFOSRDownload(capture_name=capture_name).download(save_dir) + return + + assert ( + self.capture_name in nerfosr_downloads + ), f"Capture name {self.capture_name} not found in {nerfosr_downloads.keys()}" + url = nerfosr_downloads[self.capture_name] + target_path = str(save_dir / f"NeRF-OSR/Data/{self.capture_name}") + os.makedirs(target_path, exist_ok=True) + download_path = Path(f"{target_path}.zip") + tmp_path = str(save_dir / ".temp") + shutil.rmtree(tmp_path, ignore_errors=True) + os.makedirs(tmp_path, exist_ok=True) + + os.system(f"curl -L '{url}' > {download_path}") + + # Extract the zip file + with zipfile.ZipFile(download_path, "r") as zip_ref: + zip_ref.extractall(tmp_path) + + inner_folders = os.listdir(tmp_path) + assert len(inner_folders) == 1, "There is more than one folder inside this zip file." + folder = os.path.join(tmp_path, inner_folders[0]) + shutil.rmtree(target_path) + shutil.move(folder, target_path) + shutil.rmtree(tmp_path) + os.remove(download_path) + + Commands = Union[ Annotated[BlenderDownload, tyro.conf.subcommand(name="blender")], Annotated[Sitcoms3DDownload, tyro.conf.subcommand(name="sitcoms3d")], @@ -367,6 +426,7 @@ Commands = Union[ Annotated[DNerfDownload, tyro.conf.subcommand(name="dnerf")], Annotated[PhototourismDownload, tyro.conf.subcommand(name="phototourism")], Annotated[SDFstudioDemoDownload, tyro.conf.subcommand(name="sdfstudio")], + Annotated[NeRFOSRDownload, tyro.conf.subcommand(name="nerfosr")], ] diff --git a/scripts/process_data.py b/scripts/process_data.py index 2b0647434..2eaa9e11c 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -2,7 +2,6 @@ """Processes a video or image sequence to a nerfstudio compatible dataset.""" -import json import shutil import sys import zipfile @@ -19,7 +18,6 @@ from nerfstudio.process_data import ( colmap_utils, equirect_utils, hloc_utils, - insta360_utils, metashape_utils, polycam_utils, process_data_utils, @@ -238,6 +236,7 @@ class ProcessImages: matching_method=self.matching_method, feature_type=feature_type, matcher_type=matcher_type, + refine_pixsfm=self.refine_pixsfm, ) else: CONSOLE.log("[bold red]Invalid combination of sfm_tool, feature_type, and matcher_type, exiting") @@ -457,148 +456,6 @@ class ProcessVideo: CONSOLE.rule() -@dataclass -class ProcessInsta360: - """Process Insta360 videos into a nerfstudio dataset. Currently this uses a center crop of the raw data - so data at the extreme edges of the video will be lost. - - Expects data from a 2 camera Insta360, single or >2 camera models will not work. - (tested with Insta360 One X2) - - This script does the following: - - 1. Converts the videos into images. - 2. Scales images to a specified size. - 3. Calculates the camera poses for each image using `COLMAP `_. - """ - - data: Path - """Path the data, It should be one of the 3 .insv files saved with each capture (Any work).""" - output_dir: Path - """Path to the output directory.""" - num_frames_target: int = 400 - """Target number of frames to use for the dataset, results may not be exact.""" - matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree" - """Feature matching method to use. Vocab tree is recommended for a balance of speed and - accuracy. Exhaustive is slower but more accurate. Sequential is faster but should only be used for videos.""" - num_downscales: int = 3 - """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3 - will downscale the images by 2x, 4x, and 8x.""" - skip_colmap: bool = False - """If True, skips COLMAP and generates transforms.json if possible.""" - colmap_cmd: str = "colmap" - """How to call the COLMAP executable.""" - use_sfm_depth: bool = False - """If True, export and use depth maps induced from SfM points.""" - include_depth_debug: bool = False - """If --use-sfm-depth and this flag is True, also export debug images showing SfM overlaid upon input images.""" - gpu: bool = True - """If True, use GPU.""" - verbose: bool = False - """If True, print extra logging.""" - - def main(self) -> None: - """Process video into a nerfstudio dataset.""" - install_checks.check_ffmpeg_installed() - install_checks.check_colmap_installed() - - self.output_dir.mkdir(parents=True, exist_ok=True) - image_dir = self.output_dir / "images" - image_dir.mkdir(parents=True, exist_ok=True) - - filename_back, filename_front = insta360_utils.get_insta360_filenames(self.data) - - if not filename_back.exists(): - raise FileNotFoundError(f"Could not find {filename_back}") - - ffprobe_cmd = f'ffprobe -v quiet -print_format json -show_streams -select_streams v:0 "{filename_back}"' - - ffprobe_output = process_data_utils.run_command(ffprobe_cmd) - - assert ffprobe_output is not None - ffprobe_decoded = json.loads(ffprobe_output) - - width, height = ffprobe_decoded["streams"][0]["width"], ffprobe_decoded["streams"][0]["height"] - - summary_log, num_extracted_frames = [], 0 - - if width / height == 1: - if not filename_front.exists(): - raise FileNotFoundError(f"Could not find {filename_front}") - # Convert video to images - summary_log, num_extracted_frames = insta360_utils.convert_insta360_to_images( - video_front=filename_front, - video_back=filename_back, - image_dir=image_dir, - num_frames_target=self.num_frames_target, - verbose=self.verbose, - ) - else: - summary_log, num_extracted_frames = insta360_utils.convert_insta360_single_file_to_images( - video=filename_back, - image_dir=image_dir, - num_frames_target=self.num_frames_target, - verbose=self.verbose, - ) - - # Downscale images - summary_log.append(process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose)) - - # Run Colmap - colmap_dir = self.output_dir / "colmap" - if not self.skip_colmap: - colmap_dir.mkdir(parents=True, exist_ok=True) - - colmap_utils.run_colmap( - image_dir=image_dir, - colmap_dir=colmap_dir, - camera_model=CAMERA_MODELS["fisheye"], - gpu=self.gpu, - verbose=self.verbose, - matching_method=self.matching_method, - colmap_cmd=self.colmap_cmd, - ) - - # Export depth maps - if self.use_sfm_depth: - depth_dir = self.output_dir / "depth" - depth_dir.mkdir(parents=True, exist_ok=True) - image_id_to_depth_path = colmap_utils.create_sfm_depth( - recon_dir=colmap_dir / "sparse" / "0", - output_dir=depth_dir, - include_depth_debug=self.include_depth_debug, - input_images_dir=image_dir, - verbose=self.verbose, - ) - summary_log.append( - process_data_utils.downscale_images( - depth_dir, self.num_downscales, folder_name="depths", nearest_neighbor=True, verbose=self.verbose - ) - ) - else: - image_id_to_depth_path = None - - # Save transforms.json - if (colmap_dir / "sparse" / "0" / "cameras.bin").exists(): - with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): - num_matched_frames = colmap_utils.colmap_to_json( - recon_dir=colmap_dir / "sparse" / "0", - output_dir=self.output_dir, - image_id_to_depth_path=image_id_to_depth_path, - image_rename_map=None, - ) - summary_log.append(f"Colmap matched {num_matched_frames} images") - summary_log.append(colmap_utils.get_matching_summary(num_extracted_frames, num_matched_frames)) - else: - CONSOLE.log("[bold yellow]Warning: could not find existing COLMAP results. Not generating transforms.json") - - CONSOLE.rule("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:") - - for summary in summary_log: - CONSOLE.print(summary, justify="center") - CONSOLE.rule() - - @dataclass class ProcessRecord3D: """Process Record3D data into a nerfstudio dataset. @@ -966,7 +823,6 @@ Commands = Union[ Annotated[ProcessPolycam, tyro.conf.subcommand(name="polycam")], Annotated[ProcessMetashape, tyro.conf.subcommand(name="metashape")], Annotated[ProcessRealityCapture, tyro.conf.subcommand(name="realitycapture")], - Annotated[ProcessInsta360, tyro.conf.subcommand(name="insta360")], Annotated[ProcessRecord3D, tyro.conf.subcommand(name="record3d")], ] diff --git a/tests/field_components/test_encodings.py b/tests/field_components/test_encodings.py index 271adddc9..dc59b79ff 100644 --- a/tests/field_components/test_encodings.py +++ b/tests/field_components/test_encodings.py @@ -171,6 +171,32 @@ def test_tensor_hash_encoder(): assert encoded.shape == (10, out_dim) +def test_kplane_encoder(): + """Test K-Planes encoder""" + + out_dim = 8 + + encoder = encodings.KPlanesEncoding( + resolution=(32, 32, 32, 16), + num_components=out_dim, + reduce="product", + ) + assert encoder.get_out_dim() == out_dim + + in_tensor = torch.randn((10, 4)) + encoded = encoder(in_tensor) + assert encoded.shape == (10, out_dim) + + encoder = encodings.KPlanesEncoding( + resolution=(32, 32, 32), + num_components=out_dim, + reduce="product", + ) + in_tensor = torch.randn((10, 3)) + encoded = encoder(in_tensor) + assert encoded.shape == (10, out_dim) + + if __name__ == "__main__": test_scaling_and_offset() test_nerf_encoder() diff --git a/tests/model_components/test_renderers.py b/tests/model_components/test_renderers.py index 984a45e7c..a1e098388 100644 --- a/tests/model_components/test_renderers.py +++ b/tests/model_components/test_renderers.py @@ -40,7 +40,7 @@ def test_sh_renderer(): sh_renderer = renderers.SHRenderer() rgb = sh_renderer(sh=sh, directions=directions, weights=weights) - assert torch.max(rgb) > 0.9 + assert torch.max(rgb) > 0.7 def test_acc_renderer(): diff --git a/tests/pipelines/test_vanilla_pipeline.py b/tests/pipelines/test_vanilla_pipeline.py index 69581ce27..c51038135 100644 --- a/tests/pipelines/test_vanilla_pipeline.py +++ b/tests/pipelines/test_vanilla_pipeline.py @@ -57,6 +57,9 @@ def test_load_state_dict(): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.register_parameter("param", nn.Parameter(torch.ones((3,)))) + self.register_module( + "module", Model(*args, **kwargs) + ) # make sure that non-DDP checkpoint won't be processed as DDP one def load_state_dict(self, state_dict, strict=True): """Mocked load_state_dict""" @@ -64,6 +67,7 @@ def test_load_state_dict(): was_called = True assert strict assert "param" in state_dict + assert "module.device_indicator_param" in state_dict config = VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( @@ -77,3 +81,19 @@ def test_load_state_dict(): pipeline.load_pipeline(state_dict, 0) assert was_called assert pipeline.model.param[0].item() == 2 + + # preparation for another test + state_dict["_model.param"].mul_(2) # pylint: disable=unsubscriptable-object + was_called = False + # pretends to be a DDP checkpoint + ddp_state_dict = {} + for key, value in state_dict.items(): + if key.startswith("_model."): + # replace the prefix "_model" with "_model.module" + ddp_state_dict["_model.module." + key[len("_model.") :]] = value + else: + ddp_state_dict[key] = value + # load DDP checkpoint + pipeline.load_pipeline(ddp_state_dict, 0) + assert was_called + assert pipeline.model.param[0].item() == 4 diff --git a/tests/process_data/test_process_images.py b/tests/process_data/test_process_images.py index 3e4fa82a9..1b86a5adb 100644 --- a/tests/process_data/test_process_images.py +++ b/tests/process_data/test_process_images.py @@ -12,12 +12,30 @@ from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataPars from nerfstudio.data.utils.colmap_parsing_utils import Camera from nerfstudio.data.utils.colmap_parsing_utils import Image as ColmapImage from nerfstudio.data.utils.colmap_parsing_utils import ( + qvec2rotmat, write_cameras_binary, write_images_binary, ) from scripts.process_data import ProcessImages +def random_quaternion(num_poses: int): + """ + Generates random rotation quatenion. + """ + u, v, w = np.random.uniform(size=(3, num_poses)) + quaternion = np.stack( + ( + np.sqrt(1 - u) * np.sin(2 * np.pi * v), + np.sqrt(1 - u) * np.cos(2 * np.pi * v), + np.sqrt(u) * np.sin(2 * np.pi * w), + np.sqrt(u) * np.cos(2 * np.pi * w), + ), + -1, + ) + return quaternion + + def test_process_images_skip_colmap(tmp_path: Path): """ Test ns-process-data images @@ -34,21 +52,21 @@ def test_process_images_skip_colmap(tmp_path: Path): ) frames = {} num_frames = 10 - qvecs = np.random.uniform(size=(num_frames, 4)) + qvecs = random_quaternion(num_frames) tvecs = np.random.uniform(size=(num_frames, 3)) - # original_poses = np.concatenate( - # ( - # np.concatenate( - # ( - # np.stack(list(map(qvec2rotmat, qvecs))), - # tvecs[:, :, None], - # ), - # -1, - # ), - # np.array([[[0, 0, 0, 1]]], dtype=qvecs.dtype).repeat(num_frames, 0), - # ), - # -2, - # ) + original_poses = np.concatenate( + ( + np.concatenate( + ( + np.stack(list(map(qvec2rotmat, qvecs))), + tvecs[:, :, None], + ), + -1, + ), + np.array([[[0, 0, 0, 1]]], dtype=qvecs.dtype).repeat(num_frames, 0), + ), + -2, + ) for i in range(num_frames): frames[i + 1] = ColmapImage(i + 1, qvecs[i], tvecs[i], 1, f"image_{i}.png", [], []) Image.new("RGB", (width, height)).save(tmp_path / "images" / f"image_{i}.png") @@ -77,4 +95,12 @@ def test_process_images_skip_colmap(tmp_path: Path): outputs = parser.get_dataparser_outputs("train") assert len(outputs.image_filenames) == 9 assert torch.is_tensor(outputs.dataparser_transform) - # TODO @jkulhanek: Add tests if the loaded poses are the same + + # Test if the original poses can be obtained back + gt_poses = original_poses[[0, 1, 2, 3, 4, 5, 6, 7, 9]] + dataparser_poses = outputs.transform_poses_to_original_space(outputs.cameras.camera_to_worlds, "opencv").numpy() + dataparser_poses = np.concatenate( + (dataparser_poses, np.array([[[0, 0, 0, 1]]]).repeat(len(dataparser_poses), 0)), 1 + ) + dataparser_poses = np.linalg.inv(dataparser_poses) + np.testing.assert_allclose(gt_poses, dataparser_poses, rtol=0, atol=1e-5)