mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-13 19:33:46 +08:00
190 lines
8.2 KiB
Python
190 lines
8.2 KiB
Python
"""
|
|
Lantern Model File
|
|
|
|
Currently this subclasses the Nerfacto model. Consider subclassing from the base Model.
|
|
"""
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, List, Literal, Tuple, Type
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torchmetrics.functional import structural_similarity_index_measure
|
|
from torchmetrics.image import PeakSignalNoiseRatio
|
|
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
|
|
|
|
from nerfstudio.field_components.spatial_distortions import SceneContraction
|
|
from nerfstudio.fields.density_fields import HashMLPDensityField
|
|
from nerfstudio.lantern.field import LanternNerfactoField
|
|
from nerfstudio.model_components.losses import (
|
|
MSELoss, distortion_loss, interlevel_loss, orientation_loss,
|
|
pred_normal_loss, scale_gradients_by_distance_squared)
|
|
from nerfstudio.model_components.ray_samplers import (ProposalNetworkSampler,
|
|
UniformSampler)
|
|
from nerfstudio.model_components.renderers import (AccumulationRenderer,
|
|
DepthRenderer,
|
|
NormalsRenderer,
|
|
RGBRenderer)
|
|
from nerfstudio.model_components.scene_colliders import NearFarCollider
|
|
from nerfstudio.model_components.shaders import NormalsShader
|
|
from nerfstudio.models.base_model import Model, ModelConfig # for custom Model
|
|
from nerfstudio.models.nerfacto import ( # for subclassing Nerfacto model
|
|
NerfactoModel, NerfactoModelConfig)
|
|
from nerfstudio.utils import colormaps
|
|
|
|
|
|
@dataclass
|
|
class LanternModelConfig(NerfactoModelConfig):
|
|
"""Template Model Configuration.
|
|
|
|
Add your custom model config parameters here.
|
|
"""
|
|
|
|
_target: Type = field(default_factory=lambda: LanternModel)
|
|
|
|
|
|
class LanternModel(NerfactoModel):
|
|
"""Template Model."""
|
|
|
|
config: LanternModelConfig
|
|
|
|
def populate_modules(self):
|
|
super().populate_modules()
|
|
|
|
if self.config.disable_scene_contraction:
|
|
scene_contraction = None
|
|
else:
|
|
scene_contraction = SceneContraction(order=float("inf"))
|
|
|
|
# Fields
|
|
self.field = LanternNerfactoField(
|
|
self.scene_box.aabb,
|
|
hidden_dim=self.config.hidden_dim,
|
|
num_levels=self.config.num_levels,
|
|
max_res=self.config.max_res,
|
|
base_res=self.config.base_res,
|
|
features_per_level=self.config.features_per_level,
|
|
log2_hashmap_size=self.config.log2_hashmap_size,
|
|
hidden_dim_color=self.config.hidden_dim_color,
|
|
hidden_dim_transient=self.config.hidden_dim_transient,
|
|
spatial_distortion=scene_contraction,
|
|
num_images=self.num_train_data,
|
|
use_pred_normals=self.config.predict_normals,
|
|
use_average_appearance_embedding=self.config.use_average_appearance_embedding,
|
|
appearance_embedding_dim=self.config.appearance_embed_dim,
|
|
implementation=self.config.implementation,
|
|
)
|
|
|
|
self.density_fns = []
|
|
num_prop_nets = self.config.num_proposal_iterations
|
|
# Build the proposal network(s)
|
|
self.proposal_networks = torch.nn.ModuleList()
|
|
if self.config.use_same_proposal_network:
|
|
assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed."
|
|
prop_net_args = self.config.proposal_net_args_list[0]
|
|
network = HashMLPDensityField(
|
|
self.scene_box.aabb,
|
|
spatial_distortion=scene_contraction,
|
|
**prop_net_args,
|
|
implementation=self.config.implementation,
|
|
)
|
|
self.proposal_networks.append(network)
|
|
self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)])
|
|
else:
|
|
for i in range(num_prop_nets):
|
|
prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)]
|
|
network = HashMLPDensityField(
|
|
self.scene_box.aabb,
|
|
spatial_distortion=scene_contraction,
|
|
**prop_net_args,
|
|
implementation=self.config.implementation,
|
|
)
|
|
self.proposal_networks.append(network)
|
|
self.density_fns.extend([network.density_fn for network in self.proposal_networks])
|
|
|
|
# Samplers
|
|
def update_schedule(step):
|
|
return np.clip(
|
|
np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]),
|
|
1,
|
|
self.config.proposal_update_every,
|
|
)
|
|
|
|
# Change proposal network initial sampler if uniform
|
|
initial_sampler = None # None is for piecewise as default (see ProposalNetworkSampler)
|
|
if self.config.proposal_initial_sampler == "uniform":
|
|
initial_sampler = UniformSampler(single_jitter=self.config.use_single_jitter)
|
|
|
|
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_iterations,
|
|
single_jitter=self.config.use_single_jitter,
|
|
update_sched=update_schedule,
|
|
initial_sampler=initial_sampler,
|
|
)
|
|
|
|
# Collider
|
|
self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)
|
|
|
|
# renderers
|
|
self.renderer_rgb = RGBRenderer(background_color=self.config.background_color)
|
|
self.renderer_accumulation = AccumulationRenderer()
|
|
self.renderer_depth = DepthRenderer()
|
|
self.renderer_normals = NormalsRenderer()
|
|
|
|
# shaders
|
|
self.normals_shader = NormalsShader()
|
|
|
|
# TODO losses
|
|
self.rgb_loss = MSELoss()
|
|
|
|
# metrics
|
|
self.psnr = PeakSignalNoiseRatio(data_range=1.0)
|
|
self.ssim = structural_similarity_index_measure
|
|
self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True)
|
|
|
|
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]]:
|
|
gt_rgb = batch["image"].to(self.device)
|
|
predicted_rgb = outputs["rgb"] # Blended with background (black if random background)
|
|
gt_rgb = self.renderer_rgb.blend_background(gt_rgb)
|
|
acc = colormaps.apply_colormap(outputs["accumulation"])
|
|
depth = colormaps.apply_depth_colormap(
|
|
outputs["depth"],
|
|
accumulation=outputs["accumulation"],
|
|
)
|
|
|
|
combined_rgb = torch.cat([gt_rgb, predicted_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
|
|
gt_rgb_tonemapped = torch.clamp(colormaps.apply_colormap(gt_rgb), 0, 1)
|
|
predicted_rgb_tonemapped = torch.clamp(colormaps.apply_colormap(predicted_rgb), 0, 1)
|
|
|
|
gt_rgb_tonemapped = torch.moveaxis(gt_rgb_tonemapped, -1, 0)[None, ...]
|
|
predicted_rgb_tonemapped = torch.moveaxis(predicted_rgb_tonemapped, -1, 0)[None, ...]
|
|
gt_rgb = torch.moveaxis(gt_rgb, -1, 0)[None, ...]
|
|
predicted_rgb = torch.moveaxis(predicted_rgb, -1, 0)[None, ...]
|
|
|
|
psnr = self.psnr(gt_rgb, predicted_rgb)
|
|
ssim = self.ssim(gt_rgb, predicted_rgb)
|
|
# tonemapping before computing lpips
|
|
lpips = self.lpips(gt_rgb_tonemapped, predicted_rgb_tonemapped)
|
|
|
|
# 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_iterations):
|
|
key = f"prop_depth_{i}"
|
|
prop_depth_i = colormaps.apply_depth_colormap(
|
|
outputs[key],
|
|
accumulation=outputs["accumulation"],
|
|
)
|
|
images_dict[key] = prop_depth_i
|
|
|
|
return metrics_dict, images_dict |