instant-ngp-bounded model (#1116)

* inerted a new model configuration instant-ngp-bounded

* added more remarks and small fixes

* went back to old near plane

* added the oprtion to get a file as well as directory in nerfstudio data parser

* pretify prev commit

* review fixes

* more clear explanation in README.md

* reverted readme changes

* fixed merge bug

* review changes. update config.data docs support for json files inputs

Co-authored-by: dror <you@example.com>

nerfstudio/configs/method_configs.py

@@ -21,6 +21,7 @@ from __future__ import annotations
 from typing import Dict
 
 import tyro
+from nerfacc import ContractionType
 
 from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
 from nerfstudio.configs.base_config import TrainerConfig, ViewerConfig
@@ -33,6 +34,9 @@ from nerfstudio.data.datamanagers.variable_res_datamanager import (
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
 from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig
 from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig
+from nerfstudio.data.dataparsers.instant_ngp_dataparser import (
+    InstantNGPDataParserConfig,
+)
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
 from nerfstudio.data.dataparsers.phototourism_dataparser import (
     PhototourismDataParserConfig,
@@ -52,7 +56,8 @@ from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig
 method_configs: Dict[str, ExperimentConfig] = {}
 descriptions = {
     "nerfacto": "Recommended real-time model tuned for real captures. This model will be continually updated.",
-    "instant-ngp": "Implementation of Instant-NGP. Recommended real-time model for bounded synthetic data.",
+    "instant-ngp": "Implementation of Instant-NGP. Recommended real-time model for unbounded scenes.",
+    "instant-ngp-bounded": "Implementation of Instant-NGP. Recommended for bounded real and synthetic scenes",
     "mipnerf": "High quality model for bounded scenes. (slow)",
     "semantic-nerfw": "Predicts semantic segmentations and filters out transient objects.",
     "vanilla-nerf": "Original NeRF model. (slow)",
@@ -110,6 +115,34 @@ method_configs["instant-ngp"] = ExperimentConfig(
     vis="viewer",
 )
 
+
+method_configs["instant-ngp-bounded"] = ExperimentConfig(
+    method_name="instant-ngp-bounded",
+    trainer=TrainerConfig(
+        steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True
+    ),
+    pipeline=DynamicBatchPipelineConfig(
+        datamanager=VanillaDataManagerConfig(dataparser=InstantNGPDataParserConfig(), train_num_rays_per_batch=8192),
+        model=InstantNGPModelConfig(
+            eval_num_rays_per_chunk=8192,
+            contraction_type=ContractionType.AABB,
+            render_step_size=0.001,
+            max_num_samples_per_ray=48,
+            near_plane=0.01,
+            background_color="black",
+        ),
+    ),
+    optimizers={
+        "fields": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        }
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=64000),
+    vis="viewer",
+)
+
+
 method_configs["mipnerf"] = ExperimentConfig(
     method_name="mipnerf",
     pipeline=VanillaPipelineConfig(

nerfstudio/data/dataparsers/instant_ngp_dataparser.py

@@ -20,6 +20,7 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Dict, Tuple, Type
 
+import imageio
 import numpy as np
 import torch
 from rich.console import Console
@@ -44,10 +45,10 @@ class InstantNGPDataParserConfig(DataParserConfig):
     _target: Type = field(default_factory=lambda: InstantNGP)
     """target class to instantiate"""
     data: Path = Path("data/ours/posterv2")
-    """Directory specifying location of data."""
+    """Directory or explicit json file path specifying location of data."""
     scale_factor: float = 1.0
     """How much to scale the camera origins by."""
-    scene_scale: float = 0.33
+    scene_scale: float = 0.3333
     """How much to scale the scene."""
 
 
@@ -58,16 +59,34 @@ class InstantNGP(DataParser):
     config: InstantNGPDataParserConfig
 
     def _generate_dataparser_outputs(self, split="train"):
+        if self.config.data.suffix == ".json":
+            meta = load_from_json(self.config.data)
+            data_dir = self.config.data.parent
+        else:
+            meta = load_from_json(self.config.data / "transforms.json")
+            data_dir = self.config.data
 
-        meta = load_from_json(self.config.data / "transforms.json")
         image_filenames = []
         poses = []
         num_skipped_image_filenames = 0
         for frame in meta["frames"]:
-            fname = self.config.data / Path(frame["file_path"])
-            if not fname:
+            fname = data_dir / Path(frame["file_path"])
+            # search for png file
+            if not fname.exists():
+                fname = data_dir / Path(frame["file_path"] + ".png")
+            if not fname.exists():
+                CONSOLE.log(f"couldnt find {fname} image")
                 num_skipped_image_filenames += 1
             else:
+                if "w" not in meta:
+                    img_0 = imageio.imread(fname)
+                    h, w = img_0.shape[:2]
+                    meta["w"] = w
+                    if "h" in meta:
+                        meta_h = meta["h"]
+                        assert meta_h == h, f"height of image dont not correspond metadata {h} != {meta_h}"
+                    else:
+                        meta["h"] = h
                 image_filenames.append(fname)
                 poses.append(np.array(frame["transform_matrix"]))
         if num_skipped_image_filenames >= 0:
@@ -84,12 +103,18 @@ class InstantNGP(DataParser):
         camera_to_world = torch.from_numpy(poses[:, :3])  # camera to world transform
 
         distortion_params = camera_utils.get_distortion_params(
-            k1=float(meta["k1"]), k2=float(meta["k2"]), p1=float(meta["p1"]), p2=float(meta["p2"])
+            k1=float(meta.get("k1", 0)),
+            k2=float(meta.get("k2", 0)),
+            k3=float(meta.get("k3", 0)),
+            k4=float(meta.get("k4", 0)),
+            p1=float(meta.get("p1", 0)),
+            p2=float(meta.get("p2", 0)),
         )
 
         # in x,y,z order
         # assumes that the scene is centered at the origin
-        aabb_scale = meta["aabb_scale"]
+        aabb_scale = 0.5 * meta.get("aabb_scale", 1)
+
         scene_box = SceneBox(
             aabb=torch.tensor(
                 [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32
@@ -98,14 +123,16 @@ class InstantNGP(DataParser):
 
         fl_x, fl_y = InstantNGP.get_focal_lengths(meta)
 
+        w, h = meta["w"], meta["h"]
+
         cameras = Cameras(
             fx=float(fl_x),
             fy=float(fl_y),
-            cx=float(meta["cx"]),
-            cy=float(meta["cy"]),
+            cx=float(meta.get("cx", 0.5 * w)),
+            cy=float(meta.get("cy", 0.5 * h)),
             distortion_params=distortion_params,
-            height=int(meta["h"]),
-            width=int(meta["w"]),
+            height=int(h),
+            width=int(w),
             camera_to_worlds=camera_to_world,
             camera_type=CameraType.PERSPECTIVE,
         )
@@ -140,12 +167,15 @@ class InstantNGP(DataParser):
         elif "camera_angle_x" in meta:
             fl_x = fov_to_focal_length(meta["camera_angle_x"], meta["w"])
 
-        if "fl_y" in meta:
-            fl_y = meta["fl_y"]
-        elif "y_fov" in meta:
-            fl_y = fov_to_focal_length(np.deg2rad(meta["y_fov"]), meta["h"])
-        elif "camera_angle_y" in meta:
-            fl_y = fov_to_focal_length(meta["camera_angle_y"], meta["h"])
+        if "camera_angle_y" not in meta or "y_fov" not in meta:
+            fl_y = fl_x
+        else:
+            if "fl_y" in meta:
+                fl_y = meta["fl_y"]
+            elif "y_fov" in meta:
+                fl_y = fov_to_focal_length(np.deg2rad(meta["y_fov"]), meta["h"])
+            elif "camera_angle_y" in meta:
+                fl_y = fov_to_focal_length(meta["camera_angle_y"], meta["h"])
 
         if fl_x == 0 or fl_y == 0:
             raise AttributeError("Focal length cannot be calculated from transforms.json (missing fields).")

nerfstudio/data/dataparsers/nerfstudio_dataparser.py

@@ -47,7 +47,7 @@ class NerfstudioDataParserConfig(DataParserConfig):
     _target: Type = field(default_factory=lambda: Nerfstudio)
     """target class to instantiate"""
     data: Path = Path("data/nerfstudio/poster")
-    """Directory specifying location of data."""
+    """Directory or explicit json file path specifying location of data."""
     scale_factor: float = 1.0
     """How much to scale the camera origins by."""
     downscale_factor: Optional[int] = None
@@ -74,7 +74,13 @@ class Nerfstudio(DataParser):
     def _generate_dataparser_outputs(self, split="train"):
         # pylint: disable=too-many-statements
 
-        meta = load_from_json(self.config.data / "transforms.json")
+        if self.config.data.suffix == ".json":
+            meta = load_from_json(self.config.data)
+            data_dir = self.config.data.parent
+        else:
+            meta = load_from_json(self.config.data / "transforms.json")
+            data_dir = self.config.data
+
         image_filenames = []
         mask_filenames = []
         poses = []
@@ -101,7 +107,7 @@ class Nerfstudio(DataParser):
 
         for frame in meta["frames"]:
             filepath = PurePath(frame["file_path"])
-            fname = self._get_fname(filepath)
+            fname = self._get_fname(filepath, data_dir)
             if not fname.exists():
                 num_skipped_image_filenames += 1
                 continue
@@ -140,7 +146,11 @@ class Nerfstudio(DataParser):
             poses.append(np.array(frame["transform_matrix"]))
             if "mask_path" in frame:
                 mask_filepath = PurePath(frame["mask_path"])
-                mask_fname = self._get_fname(mask_filepath, downsample_folder_prefix="masks_")
+                mask_fname = self._get_fname(
+                    mask_filepath,
+                    data_dir,
+                    downsample_folder_prefix="masks_",
+                )
                 mask_filenames.append(mask_fname)
         if num_skipped_image_filenames >= 0:
             CONSOLE.log(f"Skipping {num_skipped_image_filenames} files in dataset split {split}.")
@@ -258,21 +268,25 @@ class Nerfstudio(DataParser):
         )
         return dataparser_outputs
 
-    def _get_fname(self, filepath: PurePath, downsample_folder_prefix="images_") -> Path:
+    def _get_fname(self, filepath: PurePath, data_dir: PurePath, downsample_folder_prefix="images_") -> Path:
         """Get the filename of the image file.
         downsample_folder_prefix can be used to point to auxillary image data, e.g. masks
+
+        filepath: the base file name of the transformations.
+        data_dir: the directory of the data that contains the transform file
+        downsample_folder_prefix: prefix of the newly generated downsampled images
         """
 
         if self.downscale_factor is None:
             if self.config.downscale_factor is None:
-                test_img = Image.open(self.config.data / filepath)
+                test_img = Image.open(data_dir / filepath)
                 h, w = test_img.size
                 max_res = max(h, w)
                 df = 0
                 while True:
                     if (max_res / 2 ** (df)) < MAX_AUTO_RESOLUTION:
                         break
-                    if not (self.config.data / f"{downsample_folder_prefix}{2**(df+1)}" / filepath.name).exists():
+                    if not (data_dir / f"{downsample_folder_prefix}{2**(df+1)}" / filepath.name).exists():
                         break
                     df += 1
 
@@ -282,5 +296,5 @@ class Nerfstudio(DataParser):
                 self.downscale_factor = self.config.downscale_factor
 
         if self.downscale_factor > 1:
-            return self.config.data / f"{downsample_folder_prefix}{self.downscale_factor}" / filepath.name
-        return self.config.data / filepath
+            return data_dir / f"{downsample_folder_prefix}{self.downscale_factor}" / filepath.name
+        return data_dir / filepath

nerfstudio/model_components/ray_samplers.py

@@ -460,7 +460,8 @@ class VolumetricSampler(Sampler):
             rays_d=rays_d,
             scene_aabb=self.scene_aabb,
             grid=self.occupancy_grid,
-            sigma_fn=self.get_sigma_fn(rays_o, rays_d),
+            # this is a workaround - using density causes crash and damage quality. should be fixed
+            sigma_fn=None,  # self.get_sigma_fn(rays_o, rays_d),
             render_step_size=render_step_size,
             near_plane=near_plane,
             far_plane=far_plane,

nerfstudio/models/instant_ngp.py

@@ -28,6 +28,7 @@ from torch.nn import Parameter
 from torchmetrics import PeakSignalNoiseRatio
 from torchmetrics.functional import structural_similarity_index_measure
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from typing_extensions import Literal
 
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.engine.callbacks import (
@@ -76,8 +77,8 @@ class InstantNGPModelConfig(ModelConfig):
     """How far along ray to stop sampling."""
     use_appearance_embedding: bool = False
     """Whether to use an appearance embedding."""
-    randomize_background: bool = True
-    """Whether to randomize the background color."""
+    background_color: Literal["random", "black", "white"] = "random"
+    """The color that is given to untrained areas."""
 
 
 class NGPModel(Model):
@@ -122,7 +123,10 @@ class NGPModel(Model):
         )
 
         # renderers
-        background_color = "random" if self.config.randomize_background else colors.WHITE
+        background_color = "random"
+        if self.config.background_color in ["white", "black"]:
+            background_color = colors.COLORS_DICT[self.config.background_color]
+
         self.renderer_rgb = RGBRenderer(background_color=background_color)
         self.renderer_accumulation = AccumulationRenderer()
         self.renderer_depth = DepthRenderer(method="expected")

scripts/benchmarking/launch_train_blender.sh

@@ -62,6 +62,13 @@ timestamp=$(date "+%Y-%m-%d_%H%M%S")
 # kill all the background jobs if terminated:
 trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM EXIT
 
+dataparser="blender-data"
+trans_file=""
+if [ "$method_name" = "instant-ngp-bounded" ]; then
+    dataparser=""
+    trans_file="/transforms_train.json"
+fi
+
 for dataset in "${DATASETS[@]}"; do
     if "$single" && [ -n "${GPU_PID[$idx]+x}" ]; then
         echo "Waiting for GPU ${GPU_IDX[$idx]}"
@@ -70,7 +77,7 @@ for dataset in "${DATASETS[@]}"; do
     fi
     export CUDA_VISIBLE_DEVICES="${GPU_IDX[$idx]}"
     ns-train "${method_name}" "${method_opts[@]}" \
-             --data="data/blender/${dataset}" \
+             --data="data/blender/${dataset}${trans_file}" \
              --experiment-name="blender_${dataset}_${tag}" \
              --trainer.relative-model-dir=nerfstudio_models/ \
              --trainer.steps-per-save=1000 \
@@ -79,7 +86,7 @@ for dataset in "${DATASETS[@]}"; do
              --logging.enable-profiler=False \
              --vis "${vis}" \
              --timestamp "$timestamp" \
-             blender-data & GPU_PID[$idx]=$!
+             ${dataparser} & GPU_PID[$idx]=$!
     echo "Launched ${method_name} ${dataset} on gpu ${GPU_IDX[$idx]}, ${tag}"
     
     # update gpu

tests/test_train.py

@@ -14,7 +14,7 @@ from nerfstudio.configs.method_configs import method_configs
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
 from scripts.train import train_loop
 
-BLACKLIST = ["base", "semantic-nerfw", "instant-ngp", "nerfacto", "phototourism"]
+BLACKLIST = ["base", "semantic-nerfw", "instant-ngp", "instant-ngp-bounded", "nerfacto", "phototourism"]
 
 
 def set_reduced_config(config: ExperimentConfig):