nerfstudio/nerfactory/models/compound.py

# 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.instant_ngp_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)

        # 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

    @torch.cuda.amp.autocast()
    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)

        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):
        image = batch["image"]
        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.item()