Background MLP and SphereCollider

nerfstudio/data/datamanagers/dreamfusion_datamanager.py

@@ -241,7 +241,7 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
             radius_mean=1.5,
             radius_std=0,
             focal_range=[1, 1],
-            vertical_rotation_range=[0, 20],
+            vertical_rotation_range=[-30, 30],
             # jitter_std=0,
         )
 

nerfstudio/field_components/field_heads.py

@@ -34,6 +34,7 @@ class FieldHeadNames(Enum):
     NORMALS = "normals"
     PRED_NORMALS = "pred_normals"
     UNCERTAINTY = "uncertainty"
+    BACKGROUND_RGB = "background_rgb"
     TRANSIENT_RGB = "transient_rgb"
     TRANSIENT_DENSITY = "transient_density"
     SEMANTICS = "semantics"

nerfstudio/fields/dreamfusion_field.py

@@ -25,7 +25,7 @@ from torch import nn
 from torch.nn.parameter import Parameter
 from torchtyping import TensorType
 
-from nerfstudio.cameras.rays import RaySamples
+from nerfstudio.cameras.rays import RayBundle, RaySamples
 from nerfstudio.data.scene_box import SceneBox
 from nerfstudio.field_components.activations import trunc_exp
 from nerfstudio.field_components.embedding import Embedding
@@ -137,6 +137,18 @@ class DreamFusionField(Field):
         )
         self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims)
 
+        self.mlp_background_color = tcnn.Network(
+            n_input_dims=self.position_encoding.n_output_dims,
+            n_output_dims=3,
+            network_config={
+                "otype": "FullyFusedMLP",
+                "activation": "ReLU",
+                "output_activation": "None",
+                "n_neurons": 64,
+                "n_hidden_layers": 2,
+            },
+        )
+
         self.mlp_head = tcnn.Network(
             n_input_dims=self.geo_feat_dim,
             n_output_dims=3,
@@ -171,6 +183,17 @@ class DreamFusionField(Field):
         density = trunc_exp(density_before_activation.to(positions))
         return density, base_mlp_out
 
+    def get_background_rgb(self, ray_bundle: RayBundle):
+        """Predicts background colors at infinity."""
+        directions = get_normalized_directions(ray_bundle.directions)
+
+        outputs_shape = ray_bundle.directions.shape[:-1]
+
+        positions_flat = self.position_encoding(directions.view(-1, 3))
+        background_rgb = self.mlp_background_color(positions_flat).view(*outputs_shape, -1).to(directions)
+        return background_rgb
+
+
     def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
         assert density_embedding is not None
         outputs = {}

nerfstudio/generative/stable_diffusion.py

@@ -13,6 +13,7 @@ from transformers import CLIPTextModel, CLIPTokenizer, logging
 CONSOLE = Console(width=120)
 logging.set_verbosity_error()
 IMG_DIM = 512
+const_scale = 0.15#8215
 
 class StableDiffusion(nn.Module):
     """Stable Diffusion implementation
@@ -144,7 +145,7 @@ class StableDiffusion(nn.Module):
 
     def latents_to_img(self, latents):
 
-        latents = 1 / 0.18215 * latents
+        latents = 1 / const_scale * latents
 
         with torch.no_grad():
             imgs = self.auto_encoder.decode(latents).sample
@@ -158,7 +159,7 @@ class StableDiffusion(nn.Module):
         imgs = 2 * imgs - 1
 
         posterior = self.auto_encoder.encode(imgs).latent_dist
-        latents = posterior.sample() * 0.18215
+        latents = posterior.sample() * const_scale
 
         return latents
 

nerfstudio/model_components/scene_colliders.py

@@ -107,6 +107,60 @@ class AABBBoxCollider(SceneCollider):
         ray_bundle.fars = fars[..., None]
         return ray_bundle
 
+class SphereCollider(SceneCollider):
+    """Module for colliding rays with the scene box to compute near and far values.
+
+    Args:
+        center: center of sphere to intersect
+        redius: radius of sphere to intersect
+    """
+
+    def __init__(self, center: TensorType[3], radius: float, near_plane: float = 0.0, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.center = center
+        self.radius = radius
+        self.near_plane = near_plane
+
+    def _intersect_with_sphere(
+        self, rays_o: TensorType["num_rays":..., 3], rays_d: TensorType["num_rays":..., 3]
+    ):
+        """Returns collection of valid rays within a specified near/far bounding box along with a mask
+        specifying which rays are valid
+
+        Args:
+            rays_o: (num_rays, 3) ray origins
+            rays_d: (num_rays, 3) ray directions
+        """
+
+        a = (rays_d * rays_d).sum(dim = -1, keepdim = True)
+        self.center = self.center.to(rays_o)
+        b = 2 * (rays_o - self.center) * rays_d
+        b = b.sum(dim = -1, keepdim = True)
+        c = (rays_o - self.center) * (rays_o - self.center)
+        c = c.sum(dim = -1, keepdim = True) - self.radius ** 2
+
+        # clamp to near plane
+        near_plane = self.near_plane if self.training else 0
+        nears = (-b - torch.sqrt(torch.square(b) - 4 * a * c)) / (2 * a)
+        fars = (-b + torch.sqrt(torch.square(b) - 4 * a * c)) / (2 * a)
+
+        nears = torch.clamp(nears, min=near_plane)
+        fars = torch.maximum(fars, nears + 1e-6)
+
+        return nears, fars
+
+    def set_nears_and_fars(self, ray_bundle: RayBundle) -> RayBundle:
+        """Intersects the rays with the scene box and updates the near and far values.
+        Populates nears and fars fields and returns the ray_bundle.
+
+        Args:
+            ray_bundle: specified ray bundle to operate on
+        """
+        nears, fars = self._intersect_with_sphere(ray_bundle.origins, ray_bundle.directions)
+        ray_bundle.nears = nears
+        ray_bundle.fars = fars
+        return ray_bundle
+
 
 class NearFarCollider(SceneCollider):
     """Sets the nears and fars with fixed values.

nerfstudio/models/dreamfusion.py

@@ -50,7 +50,7 @@ from nerfstudio.model_components.renderers import (
     NormalsRenderer,
     RGBRenderer,
 )
-from nerfstudio.model_components.scene_colliders import AABBBoxCollider
+from nerfstudio.model_components.scene_colliders import AABBBoxCollider, SphereCollider
 from nerfstudio.models.base_model import Model, ModelConfig
 from nerfstudio.utils import colormaps, colors, math, misc
 
@@ -153,6 +153,7 @@ class DreamFusionModel(Model):
         # colliders
         if self.config.enable_collider:
             self.collider = AABBBoxCollider(scene_box=self.scene_box)
+        # self.collider = SphereCollider(torch.Tensor([0, 0, 0]), 1.0)
 
     def get_param_groups(self) -> Dict[str, List[Parameter]]:
         param_groups = {}
@@ -161,13 +162,15 @@ class DreamFusionModel(Model):
 
     def get_outputs(self, ray_bundle: RayBundle):
         # uniform sampling
+        background_rgb = self.field.get_background_rgb(ray_bundle)
+        # ray_bundle = self.sphere_collider(ray_bundle)
         ray_samples_uniform = self.sampler_uniform(ray_bundle)
         field_outputs = self.field(ray_samples_uniform, compute_normals=True)
 
         density = field_outputs[FieldHeadNames.DENSITY]
         if self.initialize_density:
             pos = ray_samples_uniform.frustums.get_positions()
-            density_blob = 50 * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
+            density_blob = 100 * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
             density += density_blob
 
         weights = ray_samples_uniform.get_weights(density)
@@ -176,8 +179,12 @@ class DreamFusionModel(Model):
         depth = self.renderer_depth(weights, ray_samples_uniform)
         rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
 
+        background = torch.clamp((1 - torch.nan_to_num(accumulation, nan=0.0)), min=0.0, max=1.0) * background_rgb
+
         outputs = {
             "rgb": rgb,
+            "background_rgb": background_rgb,
+            "background": background,
             "accumulation": accumulation,
             "depth": depth,
         }
@@ -192,17 +199,17 @@ class DreamFusionModel(Model):
             light_d = ray_bundle.origins[0]
         light_d = math.safe_normalize(light_d)
 
-        ratio = 0.1
+        ratio = 0.3
         lambertian = ratio + (1 - ratio) * (normals @ light_d).clamp(min=0)
-        # lambertian = ratio + (1 - ratio) * (pred_normals @ light_d).clamp(min=0)  # [N,]
 
         shaded = lambertian.unsqueeze(-1).repeat(1, 3)
         shaded_albedo = rgb * lambertian.unsqueeze(-1)
 
         outputs["normals"] = normals
         outputs["pred_normals"] = pred_normals
-        outputs["shaded"] = shaded
+        outputs["shaded"] = torch.nan_to_num(accumulation, nan=0.0) * shaded + torch.ones_like(shaded) * torch.clamp((1 - torch.nan_to_num(accumulation, nan=0.0)), min=0.0, max=1.0)
         outputs["shaded_albedo"] = shaded_albedo
+        outputs["render"] = torch.nan_to_num(accumulation, nan=0.0) * shaded_albedo + background
 
         if self.train_normals:
 

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass, field
 from typing import Type
 
+import matplotlib.pyplot as plt
+import numpy as np
 from torch.cuda.amp.grad_scaler import GradScaler
 from typing_extensions import Literal
 
@@ -58,7 +60,19 @@ class DreamfusionPipeline(VanillaPipeline):
         model_outputs = self.model(ray_bundle)
 
         # Just uses albedo for now
-        albedo_output = model_outputs["rgb"].view(1, 64, 64, 3).permute(0, 3, 1, 2)
+        albedo_output = model_outputs["render"].view(1, 64, 64, 3).permute(0, 3, 1, 2)
+
+        accumulation = model_outputs["accumulation"].view(64, 64).detach().cpu().numpy()
+        accumulation = np.clip(accumulation, 0.0, 1.0)
+        plt.imsave('nerf_accumulation.jpg', accumulation)
+
+        background = model_outputs["background"].view(64, 64, 3).detach().cpu().numpy()
+        background = np.clip(background, 0.0, 1.0)
+        plt.imsave('nerf_background.jpg', background)
+
+        shaded = model_outputs["shaded"].view(64, 64, 3).detach().cpu().numpy()
+        shaded = np.clip(shaded, 0.0, 1.0)
+        plt.imsave('nerf_textureless.jpg', shaded)
 
         sds_loss, latents, grad = self.sd.sds_loss(self.text_embedding, albedo_output)