diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index b02d782b9..3d4b694d9 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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( diff --git a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py index 289e23d55..6a3478749 100644 --- a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py +++ b/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).") diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index d2c8c2d2e..e311a08cc 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/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 diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index a65bd2128..4bc16fc05 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/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, diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index f30d2ad9a..a67c1ec7d 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/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") diff --git a/scripts/benchmarking/launch_train_blender.sh b/scripts/benchmarking/launch_train_blender.sh index e57561e02..332ec95ef 100755 --- a/scripts/benchmarking/launch_train_blender.sh +++ b/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 diff --git a/tests/test_train.py b/tests/test_train.py index dfc2d651e..f7e860275 100644 --- a/tests/test_train.py +++ b/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):