From 6f4aad451eb98e987b9b8a0250972f23bbb11c65 Mon Sep 17 00:00:00 2001
From: terrance <terrance.3512@gmail.com>
Date: Wed, 23 Nov 2022 08:43:57 +0000
Subject: [PATCH] starting the framework for dreamfusion

---
 nerfstudio/fields/dreamfusion_field.py | 215 +++++++++++++++++++++++++
 nerfstudio/models/dreamfusion.py       | 160 ++++++++++++++++++
 2 files changed, 375 insertions(+)
 create mode 100644 nerfstudio/fields/dreamfusion_field.py
 create mode 100644 nerfstudio/models/dreamfusion.py

diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py
new file mode 100644
index 000000000..88734d061
--- /dev/null
+++ b/nerfstudio/fields/dreamfusion_field.py
@@ -0,0 +1,215 @@
+# 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.
+
+"""
+Field for dreamfusion model
+"""
+
+
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn.parameter import Parameter
+from torchtyping import TensorType
+
+from nerfstudio.cameras.rays import RaySamples
+from nerfstudio.data.scene_box import SceneBox
+from nerfstudio.field_components.activations import trunc_exp
+from nerfstudio.field_components.embedding import Embedding
+from nerfstudio.field_components.encodings import Encoding, HashEncoding, SHEncoding
+from nerfstudio.field_components.field_heads import (
+    DensityFieldHead,
+    FieldHead,
+    FieldHeadNames,
+    PredNormalsFieldHead,
+    RGBFieldHead,
+)
+from nerfstudio.field_components.mlp import MLP
+from nerfstudio.field_components.spatial_distortions import (
+    SceneContraction,
+    SpatialDistortion,
+)
+from nerfstudio.fields.base_field import Field
+
+try:
+    import tinycudann as tcnn
+except ImportError:
+    # tinycudann module doesn't exist
+    pass
+
+
+def get_normalized_directions(directions: TensorType["bs":..., 3]):
+    """SH encoding must be in the range [0, 1]
+
+    Args:
+        directions: batch of directions
+    """
+    return (directions + 1.0) / 2.0
+
+
+class DreamFusionField(Field):
+    """DreamFusion Field that uses TCNN
+
+    Args:
+        aabb: parameters of scene aabb bounds
+        num_layers: number of hidden layers
+        hidden_dim: dimension of hidden layers
+        geo_feat_dim: output geo feat dimensions
+        num_layers_color: number of hidden layers for color network
+        hidden_dim_color: dimension of hidden layers for color network
+        appearance_embedding_dim: dimension of appearance embedding
+        transient_embedding_dim: dimension of transient embedding
+        use_average_appearance_embedding: whether to use average appearance embedding or zeros for inference
+        spatial_distortion: spatial distortion to apply to the scene
+    """
+
+    def __init__(
+        self,
+        aabb,
+        num_layers: int = 2,
+        hidden_dim: int = 64,
+        geo_feat_dim: int = 15,
+        num_layers_color: int = 3,
+        hidden_dim_color: int = 64,
+        hidden_dim_transient: int = 64,
+        use_pred_normals: bool = False,
+        spatial_distortion: Optional[SpatialDistortion] = None,
+    ) -> None:
+        super().__init__()
+
+        self.aabb = Parameter(aabb, requires_grad=False)
+        self.geo_feat_dim = geo_feat_dim
+
+        self.spatial_distortion = spatial_distortion
+        self.use_pred_normals = use_pred_normals
+
+        num_levels = 16
+        max_res = 1024
+        base_res = 16
+        log2_hashmap_size = 19
+        features_per_level = 2
+        growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
+
+        self.direction_encoding = tcnn.Encoding(
+            n_input_dims=3,
+            encoding_config={
+                "otype": "SphericalHarmonics",
+                "degree": 4,
+            },
+        )
+
+        self.position_encoding = tcnn.Encoding(
+            n_input_dims=3,
+            encoding_config={"otype": "Frequency", "n_frequencies": 2},
+        )
+
+        self.mlp_base = tcnn.NetworkWithInputEncoding(
+            n_input_dims=3,
+            n_output_dims=1 + self.geo_feat_dim,
+            encoding_config={
+                "otype": "HashGrid",
+                "n_levels": num_levels,
+                "n_features_per_level": features_per_level,
+                "log2_hashmap_size": log2_hashmap_size,
+                "base_resolution": base_res,
+                "per_level_scale": growth_factor,
+            },
+            network_config={
+                "otype": "FullyFusedMLP",
+                "activation": "ReLU",
+                "output_activation": "None",
+                "n_neurons": hidden_dim,
+                "n_hidden_layers": num_layers - 1,
+            },
+        )
+
+        # predicted normals
+        if self.use_pred_normals:
+            self.mlp_pred_normals = tcnn.Network(
+                n_input_dims=self.geo_feat_dim + self.position_encoding.n_output_dims,
+                n_output_dims=hidden_dim_transient,
+                network_config={
+                    "otype": "FullyFusedMLP",
+                    "activation": "ReLU",
+                    "output_activation": "None",
+                    "n_neurons": 64,
+                    "n_hidden_layers": 2,
+                },
+            )
+            self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims)
+
+        self.mlp_head = tcnn.Network(
+            n_input_dims=self.direction_encoding.n_output_dims + self.geo_feat_dim,
+            n_output_dims=3,
+            network_config={
+                "otype": "FullyFusedMLP",
+                "activation": "ReLU",
+                "output_activation": "Sigmoid",
+                "n_neurons": hidden_dim_color,
+                "n_hidden_layers": num_layers_color - 1,
+            },
+        )
+
+    def get_density(self, ray_samples: RaySamples):
+        """Computes and returns the densities."""
+        if self.spatial_distortion is not None:
+            positions = ray_samples.frustums.get_positions()
+            positions = self.spatial_distortion(positions)
+            positions = (positions + 2.0) / 4.0
+        else:
+            positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
+        self._sample_locations = positions
+        if not self._sample_locations.requires_grad:
+            self._sample_locations.requires_grad = True
+        positions_flat = positions.view(-1, 3)
+        h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1)
+        density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)
+        self._density_before_activation = density_before_activation
+
+        # Rectifying the density with an exponential is much more stable than a ReLU or
+        # softplus, because it enables high post-activation (float32) density outputs
+        # from smaller internal (float16) parameters.
+        density = trunc_exp(density_before_activation.to(positions))
+        return density, base_mlp_out
+
+    def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
+        assert density_embedding is not None
+        outputs = {}
+
+        directions = get_normalized_directions(ray_samples.frustums.directions)
+        directions_flat = directions.view(-1, 3)
+        d = self.direction_encoding(directions_flat)
+
+        outputs_shape = ray_samples.frustums.directions.shape[:-1]
+
+        # predicted normals
+        if self.use_pred_normals:
+            positions = ray_samples.frustums.get_positions()
+
+            positions_flat = self.position_encoding(positions.view(-1, 3))
+            pred_normals_inp = torch.cat([positions_flat, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
+
+            x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions)
+            outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x)
+
+        h = torch.cat(
+            [d, density_embedding.view(-1, self.geo_feat_dim)],
+            dim=-1,
+        )
+        rgb = self.mlp_head(h).view(*outputs_shape, -1).to(directions)
+        outputs.update({FieldHeadNames.RGB: rgb})
+
+        return outputs
diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py
new file mode 100644
index 000000000..4963e1dfc
--- /dev/null
+++ b/nerfstudio/models/dreamfusion.py
@@ -0,0 +1,160 @@
+# 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.
+
+"""
+DreamFusion implementation.
+"""
+
+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 nerfstudio.cameras.rays import RayBundle
+
+from nerfstudio.field_components.encodings import NeRFEncoding, 
+from nerfstudio.field_components.field_heads import FieldHeadNames
+from nerfstudio.fields.dreamfusion_field import DreamFusionField
+from nerfstudio.model_components.losses import MSELoss
+from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler
+from nerfstudio.model_components.renderers import (
+    AccumulationRenderer,
+    DepthRenderer,
+    RGBRenderer,
+)
+from nerfstudio.model_components.scene_colliders import AABBBoxCollider
+from nerfstudio.models.base_model import Model, ModelConfig
+from nerfstudio.utils import colormaps, colors, misc
+
+
+@dataclass
+class DreamFusionModelConfig(ModelConfig):
+    """DreamFusion model config"""
+
+    _target: Type = field(default_factory=lambda: DreamFusionModel)
+    """target class to instantiate"""
+
+    num_samples: int = 256
+    """Number of samples in field evaluation"""
+
+
+class DreamFusionModel(Model):
+    """DreamFusionModel Model
+
+    Args:
+        config: DreamFusion configuration to instantiate model
+    """
+
+    def __init__(
+        self,
+        config: DreamFusionModelConfig,
+        **kwargs,
+    ) -> None:
+        self.num_samples = config.num_samples
+        super().__init__(config=config, **kwargs)
+
+    def populate_modules(self):
+        """Set the fields and modules"""
+        super().populate_modules()
+
+        # setting up fields
+
+        self.field = DreamFusionField(self.scene_box.aabb)
+
+        # samplers
+        self.sampler_uniform = UniformSampler(num_samples=self.num_samples, single_jitter=True)
+        self.sampler_pdf = PDFSampler(num_samples=self.num_samples // 2, single_jitter=True)
+
+        # renderers
+        self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)
+        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()
+
+        # colliders
+        if self.config.enable_collider:
+            self.collider = AABBBoxCollider(scene_box=self.scene_box)
+
+    def get_param_groups(self) -> Dict[str, List[Parameter]]:
+        param_groups = {}
+        param_groups["fields"] = list(self.field.parameters())
+        return param_groups
+
+    def get_outputs(self, ray_bundle: RayBundle):
+        # uniform sampling
+        ray_samples_uniform = self.sampler_uniform(ray_bundle)
+        field_outputs = self.field(ray_samples_uniform)
+        weights = ray_samples_uniform.get_weights(field_outputs[FieldHeadNames.DENSITY])
+
+        accumulation = self.renderer_accumulation(weights)
+        depth = self.renderer_depth(weights, ray_samples_uniform)
+        rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
+
+        outputs = {"rgb": rgb, "accumulation": accumulation, "depth": depth}
+        return outputs
+
+    def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
+        # TODO: modify to not take in data images, need to get diffusion loss here
+        # Scaling metrics by coefficients to create the losses.
+        device = outputs["rgb"].device
+        image = batch["image"].to(device)
+
+        rgb_loss = self.rgb_loss(image, outputs["rgb"])
+
+        loss_dict = {"rgb_loss": rgb_loss}
+        loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients)
+        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(outputs["rgb"].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)
+
+        # 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)
+
+        metrics_dict = {
+            "psnr": float(psnr.item()),
+            "ssim": float(ssim.item()),
+            "lpips": float(lpips.item()),
+        }
+        images_dict = {"img": combined_rgb, "accumulation": acc, "depth": depth}
+        return metrics_dict, images_dict