diff --git a/nerfstudio/models/nerfacto_4k.py b/nerfstudio/models/nerfacto_4k.py new file mode 100644 index 000000000..c899e7a6f --- /dev/null +++ b/nerfstudio/models/nerfacto_4k.py @@ -0,0 +1,172 @@ +# 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. + +""" +NeRF implementation that combines many recent advancements, optimized for hi-res images. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Dict, Tuple + +import torch + +from nerfstudio.cameras.rays import RayBundle +from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.model_components.losses import ( + distortion_loss, + interlevel_loss, + orientation_loss, + pred_normal_loss, +) +from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig +from nerfstudio.utils import colormaps + + +@dataclass +class Nerfacto4KModelConfig(NerfactoModelConfig): + """Nerfacto-4K Model Config""" + + +class Nerfacto4KModel(NerfactoModel): + """Nerfacto model + + Args: + config: Nerfacto configuration to instantiate model + """ + + config: Nerfacto4KModelConfig + + def populate_modules(self): + """Set the fields and modules.""" + super().populate_modules() + + 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, compute_normals=self.config.predict_normals) + 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, + } + + if self.config.predict_normals: + outputs["normals"] = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) + outputs["pred_normals"] = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) + + # These use a lot of GPU memory, so we avoid storing them for eval. + if self.training: + outputs["weights_list"] = weights_list + outputs["ray_samples_list"] = ray_samples_list + + if self.training and self.config.predict_normals: + outputs["rendered_orientation_loss"] = orientation_loss( + weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions + ) + + outputs["rendered_pred_normal_loss"] = pred_normal_loss( + weights.detach(), + field_outputs[FieldHeadNames.NORMALS].detach(), + field_outputs[FieldHeadNames.PRED_NORMALS], + ) + + for i in range(self.config.num_proposal_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) + if self.training: + 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"]) + if self.training: + loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( + outputs["weights_list"], outputs["ray_samples_list"] + ) + assert metrics_dict is not None and "distortion" in metrics_dict + loss_dict["distortion_loss"] = self.config.distortion_loss_mult * metrics_dict["distortion"] + if self.config.predict_normals: + # orientation loss for computed normals + loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean( + outputs["rendered_orientation_loss"] + ) + + # ground truth supervision for normals + loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean( + outputs["rendered_pred_normal_loss"] + ) + 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 = colormaps.apply_colormap(outputs["accumulation"]) + depth = colormaps.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} + + # normals to RGB for visualization. TODO: use a colormap + if "normals" in outputs: + images_dict["normals"] = (outputs["normals"] + 1.0) / 2.0 + if "pred_normals" in outputs: + images_dict["pred_normals"] = (outputs["pred_normals"] + 1.0) / 2.0 + + 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