tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-14 03:44:02 +08:00
Code Activity

Remove mipnerf 360 (#659)

Browse Source 
remove mipnerf 360
This commit is contained in:
Matthew Tancik authored and GitHub committed 2022-10-04 16:29:11 -07:00
1 parent 70e7fdde8e
commit 654f070553
10 files changed
+3 -418
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.vscode/launch.json
Vendored
-8
View File
@@ -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",
README.md
+1 -1
View File
@@ -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.
docs/developer_guides/pipelines/creating_new_methods.ipynb
+1 -2
View File
@@ -37,8 +37,7 @@
"<summary>Models</summary>\n",
"<br>\n",
"<li>Vanilla NeRF Model</li>\n",
"<li>Mip NeRF Model</li>\n",
"<li>Mip NeRF 360 Model</li>\n",
"<li>mipNeRF Model</li>\n",
"<li>Instant NGP Model</li>\n",
"<li>NeRF-W Model</li>\n",
"<li>Semantic NeRF Model</li>\n",
docs/developer_guides/pipelines/dataparsers.md
-8
View File
@@ -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.
docs/index.md
-1
View File
@@ -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
nerfstudio/configs/method_configs.py
-25
View File
@@ -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(
nerfstudio/data/datamanagers.py
-2
View File
@@ -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(),
},
nerfstudio/data/dataparsers/mipnerf_dataparser.py
-164
View File
@@ -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
nerfstudio/models/mipnerf_360.py
-206
View File
@@ -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
tests/test_train.py
+1 -1
View File
@@ -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):