diff --git a/.vscode/launch.json b/.vscode/launch.json index 1192e3071..b3390e995 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -85,14 +85,6 @@ "--pipeline.datamanager.train-dataparser.scene-scale=4" ] }, - { - "name": "Python: train mipnerf 360", - "type": "python", - "request": "launch", - "program": "scripts/train.py", - "console": "integratedTerminal", - "args": ["mipnerf_360"] - }, { "name": "Python: train tensorf", "type": "python", diff --git a/README.md b/README.md index f4f6ed4be..f130e3403 100644 --- a/README.md +++ b/README.md @@ -139,7 +139,7 @@ ns-train nerfacto --help # Run with nerfacto model. ns-train nerfacto -# We provide support for many existing models. E.g. to run mipnerf. +# We provide support for other models. E.g. to run instant-ngp. ns-train instant-ngp # To train on your custom data. diff --git a/docs/developer_guides/pipelines/creating_new_methods.ipynb b/docs/developer_guides/pipelines/creating_new_methods.ipynb index 9d9f4d1fa..489330953 100644 --- a/docs/developer_guides/pipelines/creating_new_methods.ipynb +++ b/docs/developer_guides/pipelines/creating_new_methods.ipynb @@ -37,8 +37,7 @@ "Models\n", "
\n", "
  • Vanilla NeRF Model
    • \n", - "
  • Mip NeRF Model
    • \n", - "
  • Mip NeRF 360 Model
    • \n", + "
  • mipNeRF Model
    • \n", "
  • Instant NGP Model
    • \n", "
  • NeRF-W Model
    • \n", "
  • Semantic NeRF Model
    • \n", diff --git a/docs/developer_guides/pipelines/dataparsers.md b/docs/developer_guides/pipelines/dataparsers.md index e0fbd7e81..3a3ef6ec5 100644 --- a/docs/developer_guides/pipelines/dataparsers.md +++ b/docs/developer_guides/pipelines/dataparsers.md @@ -143,14 +143,6 @@ This supports the Instant NGP dataset. See the code! ``` -###### MipNeRF 360 - -```{button-link} https://github.com/nerfstudio-project/nerfstudio/blob/master/nerfstudio/data/dataparsers/mipnerf_dataparser.py -:color: primary -:outline: -See the code! -``` - ###### Record3D This dataparser can use recorded data from a >= iPhone 12 Pro using the Record3D app. Record a video and export with the `EXR + JPG sequence` format. Unzip export and `rgb` folder before training. diff --git a/docs/index.md b/docs/index.md index 180dfeaac..b5c474e09 100644 --- a/docs/index.md +++ b/docs/index.md @@ -82,7 +82,6 @@ This documentation is organized into 3 parts: - [**Nerfacto**](https://github.com/nerfstudio-project/nerfstudio/blob/master/nerfstudio/models/nerfacto.py): our defacto NeRF method combines modules focused on quality with modules focused on faster rendering. Nerfstudio easily lets us experiment with the best of both worlds! - [NeRF](https://www.matthewtancik.com/nerf): Representing Scenes as Neural Radiance Fields for View Synthesis - [Instant NGP](https://nvlabs.github.io/instant-ngp/): Instant Neural Graphics Primitives with a Multiresolution Hash Encoding -- [Mipnerf-360](https://jonbarron.info/mipnerf360/): Unbounded Anti-Aliased Neural Radiance Fields - [Mipnerf](https://jonbarron.info/mipnerf/): A Multiscale Representation for Anti-Aliasing Neural Radiance Fields - [NerfW](https://nerf-w.github.io/): Neural Radiance Fields for Unconstrained Photo Collections - [Semantic NeRF](https://shuaifengzhi.com/Semantic-NeRF/): In-Place Scene Labelling and Understanding with Implicit Scene Representation diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index bfc25d390..614049c6f 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -26,14 +26,12 @@ from nerfstudio.configs.base_config import Config, TrainerConfig, ViewerConfig from nerfstudio.data.datamanagers import VanillaDataManagerConfig from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig -from nerfstudio.data.dataparsers.mipnerf_dataparser import MipNerf360DataParserConfig from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig from nerfstudio.engine.schedulers import SchedulerConfig from nerfstudio.models.base_model import VanillaModelConfig from nerfstudio.models.instant_ngp import InstantNGPModelConfig from nerfstudio.models.mipnerf import MipNerfModel -from nerfstudio.models.mipnerf_360 import MipNerf360Model from nerfstudio.models.nerfacto import NerfactoModelConfig from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig from nerfstudio.models.tensorf import TensoRFModelConfig @@ -46,7 +44,6 @@ descriptions = { "nerfacto": "[bold green]Recommended[/bold green] 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.", - "mipnerf-360": "High quality model for unbounded 360 degree scenes. [red]*slow*", "mipnerf": "High quality model for bounded scenes. [red]*slow*", "semantic-nerfw": "Predicts semantic segmentations and filters out transient objects.", "vanilla-nerf": "Original NeRF model. [red]*slow*", @@ -93,28 +90,6 @@ method_configs["instant-ngp"] = Config( vis="viewer", ) -method_configs["mipnerf-360"] = Config( - method_name="mipnerf-360", - trainer=TrainerConfig(steps_per_eval_batch=200), - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig(dataparser=MipNerf360DataParserConfig(), train_num_rays_per_batch=8192), - model=VanillaModelConfig( - _target=MipNerf360Model, - collider_params={"near_plane": 0.5, "far_plane": 20.0}, - loss_coefficients={"ray_loss_coarse": 1.0, "ray_loss_fine": 1.0}, - num_coarse_samples=128, - num_importance_samples=128, - eval_num_rays_per_chunk=8192, - ), - ), - optimizers={ - "fields": { - "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), - "scheduler": None, - } - }, -) - method_configs["mipnerf"] = Config( method_name="mipnerf", pipeline=VanillaPipelineConfig( diff --git a/nerfstudio/data/datamanagers.py b/nerfstudio/data/datamanagers.py index f13a73b0a..64bd8ecd1 100644 --- a/nerfstudio/data/datamanagers.py +++ b/nerfstudio/data/datamanagers.py @@ -36,7 +36,6 @@ from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConf from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( InstantNGPDataParserConfig, ) -from nerfstudio.data.dataparsers.mipnerf_dataparser import MipNerf360DataParserConfig from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig from nerfstudio.data.pixel_samplers import PixelSampler @@ -56,7 +55,6 @@ AnnotatedDataParserUnion = dcargs.conf.OmitSubcommandPrefixes[ # Omit prefixes "nerfstudio-data": NerfstudioDataParserConfig(), "blender-data": BlenderDataParserConfig(), "friends-data": FriendsDataParserConfig(), - "mipnerf-360-data": MipNerf360DataParserConfig(), "instant-ngp-data": InstantNGPDataParserConfig(), "record3d-data": Record3DDataParserConfig(), }, diff --git a/nerfstudio/data/dataparsers/mipnerf_dataparser.py b/nerfstudio/data/dataparsers/mipnerf_dataparser.py deleted file mode 100644 index c4eb2399e..000000000 --- a/nerfstudio/data/dataparsers/mipnerf_dataparser.py +++ /dev/null @@ -1,164 +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 parser for mipnerf dataset""" -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Type - -import imageio -import numpy as np -import torch - -from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, -) -from nerfstudio.data.scene_box import SceneBox - - -@dataclass -class MipNerf360DataParserConfig(DataParserConfig): - """Mipnerf 360 dataset parser config""" - - _target: Type = field(default_factory=lambda: Mipnerf360) - """target class to instantiate""" - data: Path = Path("data/mipnerf_360/garden") - """Directory specifying location of data.""" - downscale_factor: int = 1 - """How much to downscale images.""" - val_skip: int = 8 - """1/val_skip images to use for validation.""" - auto_scale: bool = True - """Scale based on pose bounds.""" - aabb_scale: float = 4 - """Scene scale.""" - - -@dataclass -class Mipnerf360(DataParser): - """MipNeRF 360 Dataset""" - - config: MipNerf360DataParserConfig - - @classmethod - def normalize_orientation(cls, poses: np.ndarray): - """Set the _up_ direction to be in the positive Y direction. - - Args: - poses: Numpy array of poses. - """ - poses_orig = poses.copy() - bottom = np.reshape([0, 0, 0, 1.0], [1, 4]) - center = poses[:, :3, 3].mean(0) - vec2 = poses[:, :3, 2].sum(0) / np.linalg.norm(poses[:, :3, 2].sum(0)) - up = poses[:, :3, 1].sum(0) - vec0 = np.cross(up, vec2) / np.linalg.norm(np.cross(up, vec2)) - vec1 = np.cross(vec2, vec0) / np.linalg.norm(np.cross(vec2, vec0)) - c2w = np.stack([vec0, vec1, vec2, center], -1) # [3, 4] - c2w = np.concatenate([c2w[:3, :4], bottom], -2) # [4, 4] - bottom = np.tile(np.reshape(bottom, [1, 1, 4]), [poses.shape[0], 1, 1]) # [BS, 1, 4] - poses = np.concatenate([poses[:, :3, :4], bottom], -2) # [BS, 4, 4] - poses = np.linalg.inv(c2w) @ poses - poses_orig[:, :3, :4] = poses[:, :3, :4] - return poses_orig - - def _generate_dataparser_outputs(self, split="train"): - image_dir = "images" - if self.config.downscale_factor > 1: - image_dir += f"_{self.config.downscale_factor}" - image_dir = self.config.data / image_dir - if not image_dir.exists(): - raise ValueError(f"Image directory {image_dir} doesn't exist") - - valid_formats = [".jpg", ".png"] - image_filenames = [] - for f in image_dir.iterdir(): - ext = f.suffix - if ext.lower() not in valid_formats: - continue - image_filenames.append(image_dir / f) - image_filenames = sorted(image_filenames) - num_images = len(image_filenames) - - poses_data = np.load(self.config.data / "poses_bounds.npy") - poses = poses_data[:, :-2].reshape([-1, 3, 5]).astype(np.float32) - bounds = poses_data[:, -2:].transpose([1, 0]) - - if num_images != poses.shape[0]: - raise RuntimeError(f"Different number of images ({num_images}), and poses ({poses.shape[0]})") - - idx_test = np.arange(num_images)[:: self.config.val_skip] - idx_train = np.array([i for i in np.arange(num_images) if i not in idx_test]) - idx = idx_train if split == "train" else idx_test - - image_filenames = np.array(image_filenames)[idx] - poses = poses[idx] - - img_0 = imageio.imread(image_filenames[0]) - image_height, image_width = img_0.shape[:2] - - poses[:, :2, 4] = np.array([image_height, image_width]) - poses[:, 2, 4] = poses[:, 2, 4] - - # Reorder pose to match our convention - poses = np.concatenate([poses[:, :, 1:2], -poses[:, :, 0:1], poses[:, :, 2:]], axis=-1) - - # Center poses and rotate. (Compute up from average of all poses) - poses = self.normalize_orientation(poses) - - # Scale factor used in mipnerf - if self.config.auto_scale: - scale_factor = 1 / (np.min(bounds) * 0.75) - poses[:, :3, 3] *= scale_factor - bounds *= scale_factor - - # Center poses - poses[:, :3, 3] = poses[:, :3, 3] - np.mean(poses[:, :3, :], axis=0)[:, 3] - - focal_length = poses[0, -1, -1] - - cx = image_width / 2.0 - cy = image_height / 2.0 - camera_to_world = torch.from_numpy(poses[:, :3, :4]) # camera to world transform - num_cameras = len(image_filenames) - num_intrinsics_params = 3 - intrinsics = torch.ones((num_cameras, num_intrinsics_params), dtype=torch.float32) - intrinsics *= torch.tensor([cx, cy, focal_length]) - - aabb = torch.tensor([[-4, -4, -4], [4, 4, 4]], dtype=torch.float32) * self.config.aabb_scale - scene_box = SceneBox(aabb=aabb) - - cameras = Cameras( - fx=focal_length, - fy=focal_length, - cx=cx, - cy=cy, - camera_to_worlds=camera_to_world, - camera_type=CameraType.PERSPECTIVE, - ) - - cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor) - - dataparser_outputs = DataparserOutputs( - image_filenames=image_filenames, - cameras=cameras, - scene_box=scene_box, - ) - - return dataparser_outputs diff --git a/nerfstudio/models/mipnerf_360.py b/nerfstudio/models/mipnerf_360.py deleted file mode 100644 index d55a89399..000000000 --- a/nerfstudio/models/mipnerf_360.py +++ /dev/null @@ -1,206 +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. - -""" -Implementation of mip-NeRF. -""" - -from __future__ import annotations - -from typing import Dict, List, Tuple - -import torch -from torch.nn import Parameter -from torchmetrics import PeakSignalNoiseRatio -from torchmetrics.functional import structural_similarity_index_measure -from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity - -from nerfstudio.cameras.rays import RayBundle -from nerfstudio.field_components.encodings import NeRFEncoding -from nerfstudio.field_components.field_heads import FieldHeadNames -from nerfstudio.field_components.spatial_distortions import SceneContraction -from nerfstudio.fields.vanilla_nerf_field import NeRFField -from nerfstudio.model_components.losses import MSELoss, distortion_loss -from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler -from nerfstudio.model_components.renderers import ( - AccumulationRenderer, - DepthRenderer, - RGBRenderer, -) -from nerfstudio.models.base_model import Model, ModelConfig -from nerfstudio.utils import colormaps, colors, misc - - -class MipNerf360Model(Model): - """mip-NeRF model - - Args: - config: MipNerf360 configuration to instantiate model - """ - - def __init__( - self, - config: ModelConfig, - **kwargs, - ) -> None: - self.field = None - super().__init__(config=config, **kwargs) - - def populate_modules(self): - """Set the fields and modules""" - super().populate_modules() - - # setting up fields - position_encoding = NeRFEncoding( - in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True - ) - direction_encoding = NeRFEncoding( - in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True - ) - - self.field = NeRFField( - position_encoding=position_encoding, - direction_encoding=direction_encoding, - use_integrated_encoding=True, - spatial_distortion=SceneContraction(), - ) - - # samplers - self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples) - self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples, include_original=False) - - # renderers - self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) - self.renderer_accumulation = AccumulationRenderer() - self.renderer_depth = DepthRenderer() - - # losses - self.rgb_loss = MSELoss() - - # metrics - self.psnr = PeakSignalNoiseRatio(data_range=1.0) - self.ssim = structural_similarity_index_measure - self.lpips = LearnedPerceptualImagePatchSimilarity() - - def get_param_groups(self) -> Dict[str, List[Parameter]]: - param_groups = {} - if self.field is None: - raise ValueError("populate_fields() must be called before get_param_groups") - param_groups["fields"] = list(self.field.parameters()) - return param_groups - - def get_outputs(self, ray_bundle: RayBundle): - - if self.field is None: - raise ValueError("populate_fields() must be called before get_outputs") - - # uniform sampling - ray_samples_uniform = self.sampler_uniform(ray_bundle) - - # First pass: - field_outputs_coarse = self.field.forward(ray_samples_uniform) - weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) - rgb_coarse = self.renderer_rgb( - rgb=field_outputs_coarse[FieldHeadNames.RGB], - weights=weights_coarse, - ) - accumulation_coarse = self.renderer_accumulation(weights_coarse) - depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform) - ray_loss_coarse = distortion_loss(ray_samples_uniform, field_outputs_coarse[FieldHeadNames.DENSITY]) - - # pdf sampling - ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) - - # Second pass: - field_outputs_fine = self.field.forward(ray_samples_pdf) - weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) - rgb_fine = self.renderer_rgb( - rgb=field_outputs_fine[FieldHeadNames.RGB], - weights=weights_fine, - ) - accumulation_fine = self.renderer_accumulation(weights_fine) - depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf) - ray_loss_fine = distortion_loss(ray_samples_pdf, field_outputs_fine[FieldHeadNames.DENSITY]) - - outputs = { - "rgb_coarse": rgb_coarse, - "rgb_fine": rgb_fine, - "accumulation_coarse": accumulation_coarse, - "accumulation_fine": accumulation_fine, - "depth_coarse": depth_coarse, - "depth_fine": depth_fine, - "ray_loss_coarse": ray_loss_coarse, - "ray_loss_fine": ray_loss_fine, - } - return outputs - - def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: - image = batch["image"].to(self.device) - rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"]) - rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"]) - loss_dict = { - "rgb_loss_coarse": rgb_loss_coarse, - "rgb_loss_fine": rgb_loss_fine, - "ray_loss_coarse": torch.mean(outputs["ray_loss_coarse"]), - "ray_loss_fine": torch.mean(outputs["ray_loss_fine"]), - } - loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) - return loss_dict - - def get_image_metrics_and_images( - self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor] - ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]: - image = batch["image"].to(outputs["rgb_coarse"].device) - rgb_coarse = outputs["rgb_coarse"] - rgb_fine = outputs["rgb_fine"] - acc_coarse = colormaps.apply_colormap(outputs["accumulation_coarse"]) - acc_fine = colormaps.apply_colormap(outputs["accumulation_fine"]) - depth_coarse = colormaps.apply_depth_colormap( - outputs["depth_coarse"], - accumulation=outputs["accumulation_coarse"], - near_plane=self.config.collider_params["near_plane"], - far_plane=self.config.collider_params["far_plane"], - ) - depth_fine = colormaps.apply_depth_colormap( - outputs["depth_fine"], - accumulation=outputs["accumulation_fine"], - near_plane=self.config.collider_params["near_plane"], - far_plane=self.config.collider_params["far_plane"], - ) - - combined_rgb = torch.cat([image, rgb_coarse, rgb_fine], dim=1) - combined_acc = torch.cat([acc_coarse, acc_fine], dim=1) - combined_depth = torch.cat([depth_coarse, depth_fine], dim=1) - - # Switch images from [H, W, C] to [1, C, H, W] for metrics computations - image = torch.moveaxis(image, -1, 0)[None, ...] - rgb_coarse = torch.moveaxis(rgb_coarse, -1, 0)[None, ...] - rgb_fine = torch.moveaxis(rgb_fine, -1, 0)[None, ...] - - coarse_psnr = self.psnr(image, rgb_coarse) - fine_psnr = self.psnr(image, rgb_fine) - fine_ssim = self.ssim(image, rgb_fine) - fine_lpips = self.lpips(image, rgb_fine) - - metrics_dict = { - "psnr": float(fine_psnr.item()), - "coarse_psnr": float(coarse_psnr), - "fine_psnr": float(fine_psnr), - "fine_ssim": float(fine_ssim), - "fine_lpips": float(fine_lpips), - "ray_loss_coarse": float(torch.mean(outputs["ray_loss_coarse"])), - "ray_loss_fine": float(torch.mean(outputs["ray_loss_fine"])), - } - images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth} - return metrics_dict, images_dict diff --git a/tests/test_train.py b/tests/test_train.py index b9890ee9b..0816088c8 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 nerfstudio.engine.trainer import train_loop -BLACKLIST = ["base", "semantic-nerfw", "mipnerf-360", "instant-ngp", "nerfacto"] +BLACKLIST = ["base", "semantic-nerfw", "instant-ngp", "nerfacto"] def set_reduced_config(config: Config):