diff --git a/nerfstudio/field_components/field_heads.py b/nerfstudio/field_components/field_heads.py index 702cafd81..15b7822cb 100644 --- a/nerfstudio/field_components/field_heads.py +++ b/nerfstudio/field_components/field_heads.py @@ -18,6 +18,7 @@ Collection of render heads from enum import Enum from typing import Callable, Optional, Union +import torch from torch import nn from torchtyping import TensorType @@ -30,12 +31,12 @@ class FieldHeadNames(Enum): RGB = "rgb" SH = "sh" DENSITY = "density" + NORMALS = "normals" + PRED_NORMALS = "pred_normals" UNCERTAINTY = "uncertainty" TRANSIENT_RGB = "transient_rgb" TRANSIENT_DENSITY = "transient_density" SEMANTICS = "semantics" - SEMANTICS_STUFF = "semantics_stuff" - SEMANTICS_THING = "semantics_thing" class FieldHead(FieldComponent): @@ -171,7 +172,7 @@ class TransientDensityFieldHead(FieldHead): class SemanticFieldHead(FieldHead): - """Semantic stuff output + """Semantic output Args: num_classes: Number of semantic classes @@ -181,3 +182,21 @@ class SemanticFieldHead(FieldHead): def __init__(self, num_classes: int, in_dim: Optional[int] = None) -> None: super().__init__(in_dim=in_dim, out_dim=num_classes, field_head_name=FieldHeadNames.SEMANTICS, activation=None) + + +class PredNormalsFieldHead(FieldHead): + """Predicted normals output. + + Args: + in_dim: input dimension. If not defined in constructor, it must be set later. + activation: output head activation + """ + + def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Tanh()) -> None: + super().__init__(in_dim=in_dim, out_dim=3, field_head_name=FieldHeadNames.PRED_NORMALS, activation=activation) + + def forward(self, in_tensor: TensorType["bs":..., "in_dim"]) -> TensorType["bs":..., "out_dim"]: + """Needed to normalize the output into valid normals.""" + out_tensor = super().forward(in_tensor) + out_tensor = torch.nn.functional.normalize(out_tensor, dim=-1) + return out_tensor diff --git a/nerfstudio/fields/base_field.py b/nerfstudio/fields/base_field.py index 7e6821999..043a29e01 100644 --- a/nerfstudio/fields/base_field.py +++ b/nerfstudio/fields/base_field.py @@ -30,6 +30,11 @@ from nerfstudio.field_components.field_heads import FieldHeadNames class Field(nn.Module): """Base class for fields.""" + def __init__(self) -> None: + super().__init__() + self._sample_locations = None + self._density_before_activation = None + def density_fn(self, positions: TensorType["bs":..., 3]) -> TensorType["bs":..., 1]: """Returns only the density. Used primarily with the density grid. @@ -57,6 +62,24 @@ class Field(nn.Module): ray_samples: Samples locations to compute density. """ + def get_normals(self) -> TensorType[..., 3]: + """Computes and returns a tensor of normals. + + Args: + density: Tensor of densities. + """ + assert self._sample_locations is not None, "Sample locations must be set before calling get_normals." + assert self._density_before_activation is not None, "Density must be set before calling get_normals." + assert ( + self._sample_locations.shape[:-1] == self._density_before_activation.shape[:-1] + ), "Sample locations and density must have the same shape besides the last dimension." + + self._density_before_activation.backward( + gradient=torch.ones_like(self._density_before_activation), inputs=self._sample_locations, retain_graph=True + ) + normals = -torch.nn.functional.normalize(self._sample_locations.grad, dim=-1) + return normals + @abstractmethod def get_outputs( self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None @@ -68,14 +91,23 @@ class Field(nn.Module): density_embedding: Density embeddings to condition on. """ - def forward(self, ray_samples: RaySamples): + def forward(self, ray_samples: RaySamples, compute_normals: bool = False): """Evaluates the field at points along the ray. Args: ray_samples: Samples to evaluate field on. """ - density, density_embedding = self.get_density(ray_samples) - field_outputs = self.get_outputs(ray_samples, density_embedding=density_embedding) + if compute_normals: + with torch.enable_grad(): + density, density_embedding = self.get_density(ray_samples) + else: + density, density_embedding = self.get_density(ray_samples) + field_outputs = self.get_outputs(ray_samples, density_embedding=density_embedding) field_outputs[FieldHeadNames.DENSITY] = density # type: ignore + + if compute_normals: + with torch.enable_grad(): + normals = self.get_normals() + field_outputs[FieldHeadNames.NORMALS] = normals # type: ignore return field_outputs diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 79b81eb37..e3273ec9d 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -34,6 +34,7 @@ from nerfstudio.field_components.field_heads import ( DensityFieldHead, FieldHead, FieldHeadNames, + PredNormalsFieldHead, RGBFieldHead, SemanticFieldHead, TransientDensityFieldHead, @@ -95,6 +96,7 @@ class TCNNNerfactoField(Field): use_transient_embedding: bool = False, use_semantics: bool = False, num_semantic_classes: int = 100, + use_pred_normals: bool = False, use_average_appearance_embedding: bool = False, spatial_distortion: Optional[SpatialDistortion] = None, ) -> None: @@ -110,6 +112,7 @@ class TCNNNerfactoField(Field): self.use_average_appearance_embedding = use_average_appearance_embedding self.use_transient_embedding = use_transient_embedding self.use_semantics = use_semantics + self.use_pred_normals = use_pred_normals num_levels = 16 max_res = 1024 @@ -126,6 +129,11 @@ class TCNNNerfactoField(Field): }, ) + 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, @@ -182,6 +190,21 @@ class TCNNNerfactoField(Field): in_dim=self.mlp_semantics.n_output_dims, num_classes=num_semantic_classes ) + # 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 + self.appearance_embedding_dim, n_output_dims=3, @@ -202,9 +225,13 @@ class TCNNNerfactoField(Field): 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 @@ -222,6 +249,8 @@ class TCNNNerfactoField(Field): directions_flat = directions.view(-1, 3) d = self.direction_encoding(directions_flat) + outputs_shape = ray_samples.frustums.directions.shape[:-1] + # appearance if self.training: embedded_appearance = self.embedding_appearance(camera_indices) @@ -245,7 +274,7 @@ class TCNNNerfactoField(Field): ], dim=-1, ) - x = self.mlp_transient(transient_input).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions) + x = self.mlp_transient(transient_input).view(*outputs_shape, -1).to(directions) outputs[FieldHeadNames.UNCERTAINTY] = self.field_head_transient_uncertainty(x) outputs[FieldHeadNames.TRANSIENT_RGB] = self.field_head_transient_rgb(x) outputs[FieldHeadNames.TRANSIENT_DENSITY] = self.field_head_transient_density(x) @@ -259,10 +288,19 @@ class TCNNNerfactoField(Field): ], dim=-1, ) - # print(semantics_input) - x = self.mlp_semantics(semantics_input).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions) + x = self.mlp_semantics(semantics_input).view(*outputs_shape, -1).to(directions) outputs[FieldHeadNames.SEMANTICS] = self.field_head_semantics(x) + # 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, @@ -271,7 +309,7 @@ class TCNNNerfactoField(Field): ], dim=-1, ) - rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions) + rgb = self.mlp_head(h).view(*outputs_shape, -1).to(directions) outputs.update({FieldHeadNames.RGB: rgb}) return outputs @@ -342,6 +380,8 @@ class TorchNerfactoField(Field): self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None ) -> Dict[FieldHeadNames, TensorType]: + outputs_shape = ray_samples.frustums.directions.shape[:-1] + if ray_samples.camera_indices is None: raise AttributeError("Camera indices are not provided.") camera_indices = ray_samples.camera_indices.squeeze() @@ -349,7 +389,7 @@ class TorchNerfactoField(Field): embedded_appearance = self.embedding_appearance(camera_indices) else: embedded_appearance = torch.zeros( - (*ray_samples.frustums.directions.shape[:-1], self.appearance_embedding_dim), + (*outputs_shape, self.appearance_embedding_dim), device=ray_samples.frustums.directions.device, ) diff --git a/nerfstudio/fields/tensorf_field.py b/nerfstudio/fields/tensorf_field.py index 4972fa1df..6d4e0458b 100644 --- a/nerfstudio/fields/tensorf_field.py +++ b/nerfstudio/fields/tensorf_field.py @@ -111,8 +111,14 @@ class TensoRFField(Field): return rgb def forward( - self, ray_samples: RaySamples, mask: Optional[TensorType] = None, bg_color: Optional[TensorType] = None + self, + ray_samples: RaySamples, + compute_normals: bool = False, + mask: Optional[TensorType] = None, + bg_color: Optional[TensorType] = None, ): + if compute_normals is True: + raise ValueError("Surface normals are not currently supported with TensoRF") if mask is not None and bg_color is not None: base_density = torch.zeros(ray_samples.shape)[:, :, None].to(mask.device) base_rgb = bg_color.repeat(ray_samples[:, :, None].shape) diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index d06b27828..1946e4abc 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -170,3 +170,27 @@ def nerfstudio_distortion_loss( loss = loss + 1 / 3.0 * torch.sum(weights**2 * (ends - starts), dim=-2) return loss + + +def orientation_loss( + weights: TensorType["bs":..., "num_samples", 1], + normals: TensorType["bs":..., "num_samples", 3], + viewdirs: TensorType["bs":..., 3], +): + """Orientation loss proposed in Ref-NeRF. + Loss that encourages that all visible normals are facing towards the camera. + """ + w = weights + n = normals + v = viewdirs + n_dot_v = (n * v[..., None, :]).sum(axis=-1) + return (w[..., 0] * torch.fmin(torch.zeros_like(n_dot_v), n_dot_v) ** 2).sum(dim=-1) + + +def pred_normal_loss( + weights: TensorType["bs":..., "num_samples", 1], + normals: TensorType["bs":..., "num_samples", 3], + pred_normals: TensorType["bs":..., "num_samples", 3], +): + """Loss between normals calculated from density and normals from prediction network.""" + return (weights[..., 0] * (1.0 - torch.sum(normals * pred_normals, dim=-1))).sum(dim=-1) diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 21726acc5..947b0ce63 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -290,6 +290,20 @@ class SemanticRenderer(nn.Module): semantics: TensorType["bs":..., "num_samples", "num_classes"], weights: TensorType["bs":..., "num_samples", 1], ) -> TensorType["bs":..., "num_classes"]: - """_summary_""" + """Calculate semantics along the ray.""" sem = torch.sum(weights * semantics, dim=-2) return sem + + +class NormalsRenderer(nn.Module): + """Calculate normals along the ray.""" + + @classmethod + def forward( + cls, + normals: TensorType["bs":..., "num_samples", 3], + weights: TensorType["bs":..., "num_samples", 1], + ) -> TensorType["bs":..., 3]: + """Calculate normals along the ray.""" + n = torch.sum(weights * normals, dim=-2) + return n diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 6929c2fda..086e3c300 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -39,11 +39,18 @@ from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.field_components.spatial_distortions import SceneContraction from nerfstudio.fields.density_fields import HashMLPDensityField from nerfstudio.fields.nerfacto_field import TCNNNerfactoField -from nerfstudio.model_components.losses import MSELoss, distortion_loss, interlevel_loss +from nerfstudio.model_components.losses import ( + MSELoss, + distortion_loss, + interlevel_loss, + orientation_loss, + pred_normal_loss, +) from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler from nerfstudio.model_components.renderers import ( AccumulationRenderer, DepthRenderer, + NormalsRenderer, RGBRenderer, ) from nerfstudio.model_components.scene_colliders import NearFarCollider @@ -85,6 +92,10 @@ class NerfactoModelConfig(ModelConfig): """Proposal loss multiplier.""" distortion_loss_mult: float = 0.002 """Distortion loss multiplier.""" + orientation_loss_mult: float = 0.0001 + """Orientation loss multipier on computed noramls.""" + pred_normal_loss_mult: float = 0.001 + """Predicted normal loss multiplier.""" use_proposal_weight_anneal: bool = True """Whether to use proposal weight annealing.""" use_average_appearance_embedding: bool = True @@ -95,6 +106,8 @@ class NerfactoModelConfig(ModelConfig): """Max num iterations for the annealing function.""" use_single_jitter: bool = True """Whether use single jitter or not for the proposal networks.""" + predict_normals: bool = False + """Whether to predict normals or not.""" class NerfactoModel(Model): @@ -117,6 +130,7 @@ class NerfactoModel(Model): self.scene_box.aabb, spatial_distortion=scene_contraction, num_images=self.num_train_data, + use_pred_normals=self.config.predict_normals, use_average_appearance_embedding=self.config.use_average_appearance_embedding, ) @@ -162,6 +176,7 @@ class NerfactoModel(Model): self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) self.renderer_accumulation = AccumulationRenderer() self.renderer_depth = DepthRenderer() + self.renderer_normals = NormalsRenderer() # losses self.rgb_loss = MSELoss() @@ -210,7 +225,7 @@ class NerfactoModel(Model): 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) + 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) @@ -219,13 +234,32 @@ class NerfactoModel(Model): depth = self.renderer_depth(weights=weights, ray_samples=ray_samples) accumulation = self.renderer_accumulation(weights=weights) - outputs = {"rgb": rgb, "accumulation": accumulation, "depth": depth} + 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]) @@ -249,6 +283,16 @@ class NerfactoModel(Model): ) 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( @@ -281,6 +325,12 @@ class NerfactoModel(Model): 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( diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index e3fa1ee65..11e6e24bc 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -219,7 +219,9 @@ class TensoRFModel(Model): ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights) # fine field: - field_outputs_fine = self.field.forward(ray_samples_pdf, acc_mask, colors.WHITE.to(weights.device)) + field_outputs_fine = self.field.forward( + ray_samples_pdf, mask=acc_mask, bg_color=colors.WHITE.to(weights.device) + ) weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY])