diff --git a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py index 5f78951d5..9bdfc9288 100644 --- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py +++ b/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, ) diff --git a/nerfstudio/field_components/field_heads.py b/nerfstudio/field_components/field_heads.py index 15b7822cb..3207a6026 100644 --- a/nerfstudio/field_components/field_heads.py +++ b/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" diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py index f9830eeed..13fd28580 100644 --- a/nerfstudio/fields/dreamfusion_field.py +++ b/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 = {} diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 5c50b4e46..23471908f 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/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 diff --git a/nerfstudio/model_components/scene_colliders.py b/nerfstudio/model_components/scene_colliders.py index 78dcf86d2..c50b41042 100644 --- a/nerfstudio/model_components/scene_colliders.py +++ b/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. diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 6fe801200..c00b13c06 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/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: diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index f0c8cbd57..b08660d79 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/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)