mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 18:44:21 +08:00
84024a4e73
* compound model now uses compound field, passing num_images in the config, working on passing image indices through ngp sampler * renaming cuda func, adding torch compound support * removing unused imports * getting num images from dataloader instead of config * fixing pylint and adding todo * undoing base config change * adding cone angle and near plane to compound model config
199 lines
7.4 KiB
Python
199 lines
7.4 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.
|
|
|
|
""" Compound model that integrates multiple NeRF Advancements
|
|
|
|
Features:
|
|
|
|
NGP Hash encoding for fast training
|
|
NGP ray sampler for fast training
|
|
Spatial contraction for training on large scenes
|
|
Appearance conditioning
|
|
|
|
TODO:
|
|
Pose optimization
|
|
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Dict, List
|
|
|
|
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
|
|
|
|
import nerfactory.cuda as nerfactory_cuda
|
|
from nerfactory.cameras.rays import RayBundle
|
|
from nerfactory.configs import base as cfg
|
|
from nerfactory.fields.compound_field import field_implementation_to_class
|
|
from nerfactory.fields.modules.field_heads import FieldHeadNames
|
|
from nerfactory.models.base import Model
|
|
from nerfactory.models.modules.ray_sampler import NGPSpacedSampler
|
|
from nerfactory.optimizers.loss import MSELoss
|
|
from nerfactory.utils import colors, misc, visualization, writer
|
|
from nerfactory.utils.callbacks import (
|
|
TrainingCallback,
|
|
TrainingCallbackAttributes,
|
|
TrainingCallbackLocation,
|
|
)
|
|
|
|
|
|
class CompoundModel(Model):
|
|
"""Compound model
|
|
|
|
Args:
|
|
field_implementation (str): one of "torch" or "tcnn", or other fields in 'field_implementation_to_class'
|
|
kwargs: additional params to pass up to the parent class model
|
|
"""
|
|
|
|
config: cfg.CompoundModelConfig
|
|
|
|
def __init__(self, config: cfg.CompoundModelConfig, **kwargs) -> None:
|
|
assert config.field_implementation in field_implementation_to_class
|
|
self.field = None
|
|
super().__init__(config=config, **kwargs)
|
|
|
|
def get_training_callbacks(
|
|
self, training_callback_attributes: TrainingCallbackAttributes
|
|
) -> List[TrainingCallback]:
|
|
assert self.density_field is not None
|
|
return [
|
|
TrainingCallback(
|
|
where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
|
|
update_every_num_iters=self.density_field.update_every_num_iters,
|
|
func=self.density_field.update_density_grid,
|
|
kwargs={"density_eval_func": self.field.density_fn}, # type: ignore
|
|
)
|
|
]
|
|
|
|
def populate_modules(self):
|
|
"""Set the fields and modules."""
|
|
super().populate_modules()
|
|
# torch or tiny-cuda-nn version
|
|
self.field = field_implementation_to_class[self.config.field_implementation](
|
|
self.scene_bounds.aabb, self.num_train_data
|
|
)
|
|
|
|
# samplers
|
|
self.sampler = NGPSpacedSampler(num_samples=self.config.num_samples, density_field=self.density_field)
|
|
|
|
# losses
|
|
self.rgb_loss = MSELoss()
|
|
|
|
# metrics
|
|
self.psnr = PeakSignalNoiseRatio(data_range=1.0)
|
|
self.ssim = structural_similarity_index_measure
|
|
self.lpips = LearnedPerceptualImagePatchSimilarity()
|
|
|
|
# no colliders default
|
|
|
|
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):
|
|
# TODO(ruilongli)
|
|
# - train test difference
|
|
# - visualize "depth_density_grid"
|
|
num_rays = len(ray_bundle)
|
|
device = ray_bundle.origins.device
|
|
|
|
if self.field is None:
|
|
raise ValueError("populate_fields() must be called before get_outputs")
|
|
ray_samples, packed_info, t_min, t_max = self.sampler(
|
|
ray_bundle, self.field.aabb, cone_angle=self.config.cone_angle, near_plane=self.config.near_plane
|
|
)
|
|
|
|
field_outputs = self.field.forward(ray_samples)
|
|
|
|
rgbs = field_outputs[FieldHeadNames.RGB]
|
|
sigmas = field_outputs[FieldHeadNames.DENSITY]
|
|
|
|
# accumulate all the rays start from zero opacity
|
|
opacities = torch.zeros((num_rays, 1), device=device)
|
|
(
|
|
accumulated_weight,
|
|
accumulated_depth,
|
|
accumulated_color,
|
|
alive_ray_mask,
|
|
) = nerfactory_cuda.VolumeRenderer.apply(
|
|
packed_info,
|
|
ray_samples.frustums.starts,
|
|
ray_samples.frustums.ends,
|
|
sigmas.contiguous(),
|
|
rgbs.contiguous(),
|
|
opacities,
|
|
)
|
|
accumulated_depth = torch.clip(accumulated_depth, t_min[:, None], t_max[:, None])
|
|
accumulated_color = accumulated_color + colors.WHITE.to(accumulated_color) * (1.0 - accumulated_weight)
|
|
|
|
outputs = {
|
|
"rgb": accumulated_color,
|
|
"accumulation": accumulated_weight,
|
|
"depth": accumulated_depth,
|
|
"alive_ray_mask": alive_ray_mask, # the rays we kept from sampler
|
|
}
|
|
return outputs
|
|
|
|
def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients):
|
|
device = self.device
|
|
image = batch["image"].to(device)
|
|
if "alive_ray_mask" in outputs:
|
|
mask = outputs["alive_ray_mask"]
|
|
rgb_loss = self.rgb_loss(image[mask], outputs["rgb"][mask])
|
|
else:
|
|
rgb_loss = self.rgb_loss(image, outputs["rgb"])
|
|
loss_dict = {"rgb_loss": rgb_loss}
|
|
loss_dict = misc.scale_dict(loss_dict, loss_coefficients)
|
|
return loss_dict
|
|
|
|
def log_test_image_outputs(self, image_idx, step, batch, outputs):
|
|
device = self.device
|
|
image = batch["image"].to(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)
|
|
|
|
writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step)
|
|
writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step)
|
|
writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step)
|
|
|
|
# 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)
|
|
|
|
writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(psnr), step=step)
|
|
writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(ssim), step=step) # type: ignore
|
|
writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(lpips), step=step)
|
|
|
|
writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(psnr), step=step)
|
|
|
|
return {"psnr": float(psnr.item())}
|