diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py index a6ecc44fb..f9830eeed 100644 --- a/nerfstudio/fields/dreamfusion_field.py +++ b/nerfstudio/fields/dreamfusion_field.py @@ -70,9 +70,6 @@ class DreamFusionField(Field): 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 """ diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 65ee2c5cf..55868cc04 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -33,7 +33,11 @@ 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.losses import ( + MSELoss, + orientation_loss, + pred_normal_loss, +) from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler from nerfstudio.model_components.renderers import ( AccumulationRenderer, @@ -55,6 +59,10 @@ class DreamFusionModelConfig(ModelConfig): num_samples: int = 256 """Number of samples in field evaluation""" + orientation_loss_mult: float = 0.0001 + """Orientation loss multipier on computed normals.""" + pred_normal_loss_mult: float = 0.001 + """Predicted normal loss multiplier.""" class DreamFusionModel(Model): @@ -64,6 +72,8 @@ class DreamFusionModel(Model): config: DreamFusion configuration to instantiate model """ + config: DreamFusionModelConfig + def __init__( self, config: DreamFusionModelConfig, @@ -120,10 +130,10 @@ class DreamFusionModel(Model): pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) # lambertian shading - light_d = ray_bundle.origins[0] # + torch.randn(3, dtype=torch.float).to(normals) + light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals) light_d = math.safe_normalize(light_d) - ratio = 0.5 + ratio = 0.1 lambertian = ratio + (1 - ratio) * (pred_normals @ light_d).clamp(min=0) # [N,] shaded = lambertian.unsqueeze(-1).repeat(1, 3) @@ -139,6 +149,17 @@ class DreamFusionModel(Model): "shaded_albedo": shaded_albedo, } + if self.training: + 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], + ) + return outputs def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: @@ -151,6 +172,16 @@ class DreamFusionModel(Model): loss_dict = {"rgb_loss": rgb_loss} loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) + if self.training: + # 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( diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index b1e772b03..15d139054 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -99,7 +99,7 @@ class NerfactoModelConfig(ModelConfig): distortion_loss_mult: float = 0.002 """Distortion loss multiplier.""" orientation_loss_mult: float = 0.0001 - """Orientation loss multipier on computed noramls.""" + """Orientation loss multipier on computed normals.""" pred_normal_loss_mult: float = 0.001 """Predicted normal loss multiplier.""" use_proposal_weight_anneal: bool = True