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

remove tensorf model (#660)

Browse Source 
* remove tensorf model

* lint
This commit is contained in:
Matthew Tancik authored and GitHub committed 2022-10-04 16:39:47 -07:00
1 parent 654f070553
commit 21f58252c1
6 files changed
+5 -299
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.vscode/launch.json
Vendored
+1 -17
View File
@@ -86,23 +86,7 @@
]
},
{
"name": "Python: train tensorf",
"type": "python",
"request": "launch",
"program": "scripts/train.py",
"console": "integratedTerminal",
"args": ["tensorf"]
},
{
"name": "Python: train nerfw",
"type": "python",
"request": "launch",
"program": "scripts/train.py",
"console": "integratedTerminal",
"args": ["nerfw"]
},
{
"name": "Python: train semantic_nerf",
"name": "Python: train semantic_nerfw",
"type": "python",
"request": "launch",
"program": "scripts/train.py",
docs/developer_guides/pipelines/creating_new_methods.ipynb
+4 -4
View File
@@ -12,7 +12,7 @@
"\n",
"Here we explain how to create custom pipelines with nerfstudio. Pipelines are composed of two components, namely a DataManager and a Model. The pipeline is responsible for routing RayBundle instances returned from the DataManager into the Model.\n",
"\n",
"A NeRF method is implemented as a Pipeline. Users can create a Pipeline with a DataManager and Model. We provide default implementations for these components prefixed with the name Vanilla (e.g., `VanillaPipeline`, `VanillaDataManager` which inherit from the `Pipeline` and `DataManager` class, respectively) and specific Models (e.g., `InstantNGPModel` or `TensoRFModel` which inherit from the `Model` class). With most NeRF methods, `VanillaPipeline` and `VanillaDataManager` can be used and only the model needs to be changed. We provide many already implmented Models for people to get started with 3D reconstruction. However, in this tutorial we will demonstrate how to customize each of the Pipeline, DataManager, and Model with an example.\n",
"A NeRF method is implemented as a Pipeline. Users can create a Pipeline with a DataManager and Model. We provide default implementations for these components prefixed with the name Vanilla (e.g., `VanillaPipeline`, `VanillaDataManager` which inherit from the `Pipeline` and `DataManager` class, respectively) and specific Models (e.g., `InstantNGPModel` which inheritw from the `Model` class). With most NeRF methods, `VanillaPipeline` and `VanillaDataManager` can be used and only the model needs to be changed. We provide many already implmented Models for people to get started with 3D reconstruction. However, in this tutorial we will demonstrate how to customize each of the Pipeline, DataManager, and Model with an example.\n",
"\n",
"We'll show how to make an incremental DataManager where we progressively add cameras to train with from multiple scenes and importance sample rays from pixels with a high loss. The extensible features that we show in this tutorial are the following:\n",
"- a DataManager that incrementally adds cameras\n",
@@ -324,7 +324,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "Python 3.8.13 ('nerfactory')",
"display_name": "Python 3.8.12 ('nerfactory')",
"language": "python",
"name": "python3"
},
@@ -338,11 +338,11 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.13"
"version": "3.8.12"
},
"vscode": {
"interpreter": {
"hash": "b7b86c17c0dfd6b313529972d424f2e075cc198a46cda005757126afe8133c7b"
"hash": "34c28001ff35fb390494047002768a8182dcf55b1b11415165e62ea61557ab83"
}
}
},
docs/index.md
-1
View File
@@ -85,7 +85,6 @@ This documentation is organized into 3 parts:
- [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
- [TensoRF](https://apchenstu.github.io/TensoRF/): Tensorial Radiance Fields
We'll be constantly growing this list! So make sure to check back in to see our updates.
docs/reference/api/models.rst
-7
View File
@@ -24,13 +24,6 @@ Semantic NeRF-W
:members:
:show-inheritance:
TensoRF
----------------
.. automodule:: nerfstudio.models.tensorf
:members:
:show-inheritance:
NeRF
----------------
nerfstudio/configs/method_configs.py
-28
View File
@@ -28,13 +28,11 @@ from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConf
from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig
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.nerfacto import NerfactoModelConfig
from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig
from nerfstudio.models.tensorf import TensoRFModelConfig
from nerfstudio.models.vanilla_nerf import NeRFModel
from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig
from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig
@@ -47,7 +45,6 @@ descriptions = {
"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*",
"tensorf": "Fast model designed for bounded scenes.",
}
method_configs["nerfacto"] = Config(
@@ -149,31 +146,6 @@ method_configs["vanilla-nerf"] = Config(
},
)
method_configs["tensorf"] = Config(
method_name="tensorf",
trainer=TrainerConfig(mixed_precision=True),
pipeline=VanillaPipelineConfig(
datamanager=VanillaDataManagerConfig(
dataparser=BlenderDataParserConfig(),
),
model=TensoRFModelConfig(),
),
optimizers={
"fields": {
"optimizer": RAdamOptimizerConfig(lr=0.001),
"scheduler": SchedulerConfig(lr_final=0.00005, max_steps=15000),
},
"position_encoding": {
"optimizer": RAdamOptimizerConfig(lr=0.02),
"scheduler": SchedulerConfig(lr_final=0.005, max_steps=15000),
},
"direction_encoding": {
"optimizer": RAdamOptimizerConfig(lr=0.02),
"scheduler": SchedulerConfig(lr_final=0.005, max_steps=15000),
},
},
)
AnnotatedBaseConfigUnion = dcargs.extras.subcommand_type_from_defaults(
defaults=method_configs, descriptions=descriptions
nerfstudio/models/tensorf.py
-242
View File
@@ -1,242 +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.
"""
TensorRF implementation.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Dict, List, Tuple, Type
import numpy as np
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.configs.config_utils import to_immutable_dict
from nerfstudio.engine.callbacks import (
TrainingCallback,
TrainingCallbackAttributes,
TrainingCallbackLocation,
)
from nerfstudio.engine.optimizers import Optimizers
from nerfstudio.field_components.encodings import TensorVMEncoding
from nerfstudio.field_components.field_heads import FieldHeadNames
from nerfstudio.fields.vanilla_nerf_field import NeRFField
from nerfstudio.model_components.losses import L1Loss, MSELoss
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, VanillaModelConfig
from nerfstudio.utils import colormaps, colors, misc
@dataclass
class TensoRFModelConfig(VanillaModelConfig):
"""TensoRF model config"""
_target: Type = field(default_factory=lambda: TensoRFModel)
"""target class to instantiate"""
init_resolution: int = 128
"""initial render resolution"""
final_resolution: int = 200
"""final render resolution"""
upsampling_iters: Tuple[int, ...] = (5000, 5500, 7000)
"""specifies a list of iteration step numbers to perform upsampling"""
loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss_coarse": 1.0, "feature_loss": 8e-5})
"""Loss specific weights."""
class TensoRFModel(Model):
"""TensoRF Model
Args:
config: TensoRF configuration to instantiate model
"""
def __init__(
self,
config: TensoRFModelConfig,
**kwargs,
) -> None:
self.init_resolution = config.init_resolution
self.upsampling_iters = config.upsampling_iters
self.upsampling_steps = (
np.round(
np.exp(
np.linspace(
np.log(config.init_resolution),
np.log(config.final_resolution),
len(config.upsampling_iters) + 1,
)
)
)
.astype("int")
.tolist()[1:]
)
super().__init__(config=config, **kwargs)
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
):
resolution = self.upsampling_steps.pop(0)
# upsample the position and direction grids
# TODO(ethan): ask Brent how to get typing to work on this... the Encoding base class type
# in NeRFField is causing the issue
self.field.position_encoding.upsample_grid(resolution)
self.field.direction_encoding.upsample_grid(resolution)
# reinitialize the optimizer
optimizers_config = training_callback_attributes.optimizers.config
training_callback_attributes.optimizers = Optimizers(
optimizers_config, training_callback_attributes.pipeline.get_param_groups()
)
# TODO(ethan): do something with the learning rate
# we don't want to reinitialize the learning rate each time
callbacks = [
TrainingCallback(
where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
iters=self.upsampling_iters,
func=reinitialize_optimizer,
args=[self, training_callback_attributes],
)
]
return callbacks
def populate_modules(self):
"""Set the fields and modules"""
super().populate_modules()
# setting up fields
position_encoding = TensorVMEncoding(
resolution=self.init_resolution,
num_components=96,
)
direction_encoding = TensorVMEncoding(
resolution=self.init_resolution,
num_components=96,
)
self.field = NeRFField(
position_encoding=position_encoding,
direction_encoding=direction_encoding,
base_mlp_num_layers=2,
base_mlp_layer_width=128,
)
# samplers
self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples)
self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples)
# renderers
self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)
self.renderer_accumulation = AccumulationRenderer()
self.renderer_depth = DepthRenderer()
# losses
self.rgb_loss = MSELoss()
self.feature_loss = L1Loss()
# 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 = {}
param_groups["fields"] = list(self.field.parameters())
param_groups["position_encoding"] = list(self.field.position_encoding.parameters())
param_groups["direction_encoding"] = list(self.field.direction_encoding.parameters())
return param_groups
def get_outputs(self, ray_bundle: RayBundle):
# uniform sampling
ray_samples_uniform = self.sampler_uniform(ray_bundle)
# coarse field:
field_outputs_coarse = self.field.forward(ray_samples_uniform)
weights = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY])
rgb = self.renderer_rgb(
rgb=field_outputs_coarse[FieldHeadNames.RGB],
weights=weights,
)
accumulation = self.renderer_accumulation(weights)
depth = self.renderer_depth(weights, ray_samples_uniform)
outputs = {"rgb": rgb, "accumulation": accumulation, "depth": depth}
return outputs
def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
# Scaling metrics by coefficients to create the losses.
device = outputs["rgb"].device
image = batch["image"].to(device)
assert isinstance(self.field.position_encoding, TensorVMEncoding)
rgb_loss = self.rgb_loss(image, outputs["rgb"])
plane_coef = self.field.position_encoding.plane_coef
line_coef = self.field.position_encoding.line_coef
plane_feature_loss = self.feature_loss(plane_coef, torch.zeros_like(plane_coef))
line_feature_loss = self.feature_loss(line_coef, torch.zeros_like(line_coef))
loss_dict = {"rgb_loss": rgb_loss, "feature_loss": plane_feature_loss + line_feature_loss}
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"].device)
rgb = outputs["rgb"]
acc = colormaps.apply_colormap(outputs["accumulation"])
depth = colormaps.apply_depth_colormap(
outputs["depth"],
accumulation=outputs["accumulation"],
near_plane=self.config.collider_params["near_plane"],
far_plane=self.config.collider_params["far_plane"],
)
combined_rgb = torch.cat([image, rgb], dim=1)
# Switch images from [H, W, C] to [1, C, H, W] for metrics computations
image = torch.moveaxis(image, -1, 0)[None, ...]
rgb = torch.moveaxis(rgb, -1, 0)[None, ...]
psnr = self.psnr(image, rgb)
ssim = self.ssim(image, rgb)
lpips = self.lpips(image, rgb)
metrics_dict = {
"psnr": float(psnr.item()),
"ssim": float(ssim.item()),
"lpips": float(lpips.item()),
}
images_dict = {"img": combined_rgb, "accumulation": acc, "depth": depth}
return metrics_dict, images_dict