mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 18:44:21 +08:00
259 lines
11 KiB
Python
259 lines
11 KiB
Python
# Copyright 2022 The Plenoptix 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 Instant NGP.
|
|
"""
|
|
|
|
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 nerfactory.cameras.rays import RayBundle
|
|
from nerfactory.fields.compound_field import TCNNCompoundField
|
|
from nerfactory.fields.density_field import DensityField
|
|
from nerfactory.fields.modules.field_heads import FieldHeadNames
|
|
from nerfactory.fields.modules.spatial_distortions import SceneContraction
|
|
from nerfactory.models.base import Model, ModelConfig
|
|
from nerfactory.models.modules.ray_sampler import ProposalNetworkSampler
|
|
from nerfactory.models.modules.scene_colliders import NearFarCollider
|
|
from nerfactory.optimizers.loss import MSELoss, distortion_loss, interlevel_loss
|
|
from nerfactory.renderers.renderers import (
|
|
AccumulationRenderer,
|
|
DepthRenderer,
|
|
RGBRenderer,
|
|
)
|
|
from nerfactory.utils import colors, visualization
|
|
from nerfactory.utils.callbacks import (
|
|
TrainingCallback,
|
|
TrainingCallbackAttributes,
|
|
TrainingCallbackLocation,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class ProposalModelConfig(ModelConfig):
|
|
"""Proposal Model Config"""
|
|
|
|
_target: Type = field(default_factory=lambda: ProposalModel)
|
|
near_plane: float = 0.05
|
|
"""How far along the ray to start sampling."""
|
|
far_plane: float = 1000.0
|
|
"""How far along the ray to stop sampling."""
|
|
randomize_background: bool = True
|
|
"""Whether to randomize the background color."""
|
|
num_proposal_samples_per_ray: Tuple[int] = (64,)
|
|
"""Number of samples per ray for the proposal network."""
|
|
num_nerf_samples_per_ray: int = 64
|
|
"""Number of samples per ray for the nerf network."""
|
|
num_proposal_network_iterations: int = 1
|
|
"""Number of proposal network iterations."""
|
|
use_same_proposal_network: bool = False
|
|
"""Use the same proposal network. Otherwise use different ones."""
|
|
interlevel_loss_mult: float = 1.0
|
|
"""Proposal loss multiplier."""
|
|
distortion_loss_mult: float = 0.01
|
|
"""Distortion loss multiplier."""
|
|
use_appearance_conditioning: bool = True
|
|
"""Whether to use appearance conditioning."""
|
|
use_proposal_weight_anneal: bool = True
|
|
"""Whether to use proposal weight annealing."""
|
|
use_average_appearance_embedding: bool = True
|
|
"""Whether to use average appearance embedding or zeros for inference."""
|
|
proposal_weights_anneal_slope: float = 10.0
|
|
"""Slope of the annealing function for the proposal weights."""
|
|
proposal_weights_anneal_max_num_iters: int = 1000
|
|
"""Max num iterations for the annealing function."""
|
|
use_single_jitter: bool = False
|
|
"""Whether use single jitter or not for first proposal network."""
|
|
|
|
|
|
class ProposalModel(Model):
|
|
"""Instant NGP model
|
|
|
|
Args:
|
|
config: instant NGP configuration to instantiate model
|
|
"""
|
|
|
|
config: ProposalModelConfig
|
|
|
|
def populate_modules(self):
|
|
"""Set the fields and modules."""
|
|
super().populate_modules()
|
|
|
|
scene_contraction = SceneContraction(order=float("inf"))
|
|
|
|
# Fields
|
|
if self.config.use_appearance_conditioning:
|
|
self.field = TCNNCompoundField(
|
|
self.scene_bounds.aabb,
|
|
spatial_distortion=scene_contraction,
|
|
num_images=self.num_train_data,
|
|
use_average_appearance_embedding=self.config.use_average_appearance_embedding,
|
|
)
|
|
else:
|
|
raise NotImplementedError("Only appearance conditioning is supported.")
|
|
|
|
# Build the proposal network(s)
|
|
self.proposal_networks = torch.nn.ModuleList()
|
|
if self.config.use_same_proposal_network:
|
|
network = DensityField(self.scene_bounds.aabb, spatial_distortion=scene_contraction)
|
|
self.proposal_networks.append(network)
|
|
self.density_fns = [network.density_fn for _ in range(self.config.num_proposal_network_iterations)]
|
|
else:
|
|
for _ in range(self.config.num_proposal_network_iterations):
|
|
network = DensityField(self.scene_bounds.aabb, spatial_distortion=scene_contraction)
|
|
self.proposal_networks.append(network)
|
|
self.density_fns = [network.density_fn for network in self.proposal_networks]
|
|
|
|
# Collider
|
|
self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)
|
|
|
|
# Samplers
|
|
self.proposal_sampler = ProposalNetworkSampler(
|
|
num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray,
|
|
num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray,
|
|
num_proposal_network_iterations=self.config.num_proposal_network_iterations,
|
|
single_jitter=self.config.use_single_jitter,
|
|
)
|
|
|
|
# renderers
|
|
background_color = None if self.config.randomize_background else colors.WHITE
|
|
self.renderer_rgb = RGBRenderer(background_color=background_color)
|
|
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 = {}
|
|
param_groups["proposal_networks"] = list(self.proposal_networks.parameters())
|
|
param_groups["fields"] = list(self.field.parameters())
|
|
return param_groups
|
|
|
|
def get_training_callbacks(
|
|
self, training_callback_attributes: TrainingCallbackAttributes
|
|
) -> List[TrainingCallback]:
|
|
callbacks = []
|
|
if self.config.use_proposal_weight_anneal:
|
|
# anneal the weights of the proposal network before doing PDF sampling
|
|
N = self.config.proposal_weights_anneal_max_num_iters
|
|
|
|
def set_anneal(step):
|
|
# https://arxiv.org/pdf/2111.12077.pdf eq. 18
|
|
train_frac = np.clip(step / N, 0, 1)
|
|
bias = lambda x, b: (b * x) / ((b - 1) * x + 1)
|
|
anneal = bias(train_frac, self.config.proposal_weights_anneal_slope)
|
|
self.proposal_sampler.set_anneal(anneal)
|
|
|
|
callbacks.append(
|
|
TrainingCallback(
|
|
where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],
|
|
update_every_num_iters=1,
|
|
func=set_anneal,
|
|
)
|
|
)
|
|
return callbacks
|
|
|
|
def get_outputs(self, ray_bundle: RayBundle):
|
|
ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)
|
|
field_outputs = self.field(ray_samples)
|
|
weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])
|
|
weights_list.append(weights)
|
|
ray_samples_list.append(ray_samples)
|
|
|
|
rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
|
|
depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)
|
|
accumulation = self.renderer_accumulation(weights=weights)
|
|
|
|
outputs = {"rgb": rgb, "accumulation": accumulation, "depth": depth}
|
|
outputs["weights_list"] = weights_list
|
|
outputs["ray_samples_list"] = ray_samples_list
|
|
|
|
for i in range(self.config.num_proposal_network_iterations):
|
|
outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])
|
|
|
|
return outputs
|
|
|
|
def get_metrics_dict(self, outputs, batch):
|
|
metrics_dict = {}
|
|
image = batch["image"].to(self.device)
|
|
metrics_dict["psnr"] = self.psnr(outputs["rgb"], image)
|
|
metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"])
|
|
return metrics_dict
|
|
|
|
def get_loss_dict(self, outputs, batch, metrics_dict=None):
|
|
loss_dict = {}
|
|
image = batch["image"].to(self.device)
|
|
loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"])
|
|
loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
|
|
outputs["weights_list"], outputs["ray_samples_list"]
|
|
)
|
|
loss_dict["distortion_loss"] = self.config.distortion_loss_mult * metrics_dict["distortion"]
|
|
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(self.device)
|
|
rgb = outputs["rgb"]
|
|
acc = visualization.apply_colormap(outputs["accumulation"])
|
|
depth = visualization.apply_depth_colormap(
|
|
outputs["depth"],
|
|
accumulation=outputs["accumulation"],
|
|
)
|
|
|
|
combined_rgb = torch.cat([image, rgb], dim=1)
|
|
combined_acc = torch.cat([acc], dim=1)
|
|
combined_depth = torch.cat([depth], 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)
|
|
|
|
# all of these metrics will be logged as scalars
|
|
metrics_dict = {"psnr": float(psnr.item()), "ssim": float(ssim)} # type: ignore
|
|
metrics_dict["lpips"] = float(lpips)
|
|
|
|
images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth}
|
|
|
|
for i in range(self.config.num_proposal_network_iterations):
|
|
key = f"prop_depth_{i}"
|
|
prop_depth_i = visualization.apply_depth_colormap(
|
|
outputs[key],
|
|
accumulation=outputs["accumulation"],
|
|
)
|
|
images_dict[key] = prop_depth_i
|
|
|
|
return metrics_dict, images_dict
|