use our distortion loss (#543)

swap out with our distortion loss

nerfactory/models/modules/ray_losses.py

@@ -21,7 +21,10 @@ from nerfactory.cameras.rays import RaySamples
 
 
 def distortion_loss(
-    ray_samples: RaySamples, densities: TensorType["bs":..., "num_samples", 1]
+    ray_samples: RaySamples,
+    densities: TensorType["bs":..., "num_samples", 1] = None,
+    weights: TensorType["bs":..., "num_samples", 1] = None,
+    scale_factor: float = 1.0,
 ) -> TensorType["bs":..., 1]:
     """Ray baserd distortion loss proposed in MipNeRF-360. Returns distortion Loss.
 
@@ -36,17 +39,25 @@ def distortion_loss(
     Args:
         ray_samples: Ray samples to compute loss over
         densities: Predicted sample densities
+        weights: Predicted weights from densities and sample locations
+        scale_factor: Scale the starts and ends by this factor
     """
+    if torch.is_tensor(densities):
+        assert not torch.is_tensor(weights), "Cannot use both densities and weights"
+        # Compute the weight at each sample location
+        weights = ray_samples.get_weights(densities)
+    if torch.is_tensor(weights):
+        assert not torch.is_tensor(densities), "Cannot use both densities and weights"
 
-    # Compute the weight at each sample location
-    weights = ray_samples.get_weights(densities)
+    starts = ray_samples.frustums.starts * scale_factor
+    ends = ray_samples.frustums.ends * scale_factor
 
-    midpoints = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2.0  # (..., num_samples, 1)
+    midpoints = (starts + ends) / 2.0  # (..., num_samples, 1)
 
     loss = (
         weights * weights[..., None, :, 0] * torch.abs(midpoints - midpoints[..., None, :, 0])
     )  # (..., num_samples, num_samples)
     loss = torch.sum(loss, dim=(-1, -2))[..., None]  # (..., num_samples)
-    loss = loss + 1 / 3.0 * torch.sum(weights**2 * (ray_samples.frustums.ends - ray_samples.frustums.starts), dim=-2)
+    loss = loss + 1 / 3.0 * torch.sum(weights**2 * (ends - starts), dim=-2)
 
     return loss

nerfactory/models/proposal.py

@@ -34,9 +34,10 @@ from nerfactory.fields.instant_ngp_field import TCNNInstantNGPField
 from nerfactory.fields.modules.field_heads import FieldHeadNames
 from nerfactory.fields.modules.spatial_distortions import SceneContraction
 from nerfactory.models.base import Model, ModelConfig
+from nerfactory.models.modules.ray_losses import distortion_loss
 from nerfactory.models.modules.ray_sampler import ProposalNetworkSampler
 from nerfactory.models.modules.scene_colliders import NearFarCollider
-from nerfactory.optimizers.loss import MSELoss, distortion_loss, interlevel_loss
+from nerfactory.optimizers.loss import MSELoss, interlevel_loss
 from nerfactory.renderers.renderers import (
     AccumulationRenderer,
     DepthRenderer,
@@ -155,11 +156,12 @@ class ProposalModel(Model):
             self.config.near_plane,
             self.config.far_plane,
         )
-        loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss(
-            outputs["weights_list"],
-            outputs["ray_samples_list"],
-            self.config.near_plane,
-            self.config.far_plane,
+        loss_dict["distortion_loss"] = self.config.distortion_loss_mult * torch.mean(
+            distortion_loss(
+                outputs["ray_samples_list"][-1],
+                weights=outputs["weights_list"][-1],
+                scale_factor=1.0 / (self.config.far_plane - self.config.near_plane),
+            )
         )
         return loss_dict
 
@@ -188,8 +190,8 @@ class ProposalModel(Model):
         lpips = self.lpips(image, rgb)
 
         # all of these metrics will be logged as scalars
-        metrics_dict = {"psnr": float(psnr.item()), "ssim": float(ssim), "lpips": float(lpips)}  # type: ignore
-        # TODO(ethan): return an image dictionary
+        metrics_dict = {"psnr": float(psnr.item()), "ssim": float(ssim)}  # type: ignore
+        metrics_dict["lpips"] = float(lpips)
 
         images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth}
 

nerfactory/optimizers/loss.py

@@ -70,20 +70,6 @@ def lossfun_outer(t, w, t_env, w_env):
     return torch.clip(w - w_outer, min=0) ** 2 / (w + EPS)
 
 
-def lossfun_distortion(t, w):
-    """
-    https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L142
-    https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L266
-    """
-    ut = (t[..., 1:] + t[..., :-1]) / 2
-    dut = torch.abs(ut[..., :, None] - ut[..., None, :])
-    loss_inter = torch.sum(w * torch.sum(w[..., None, :] * dut, dim=-1), dim=-1)
-
-    loss_intra = torch.sum(w**2 * (t[..., 1:] - t[..., :-1]), dim=-1) / 3
-
-    return loss_inter + loss_intra
-
-
 def ray_samples_to_sdist(ray_samples, near_plane, far_plane):
     """Convert ray samples to s space"""
     starts = (ray_samples.frustums.starts - near_plane) / (far_plane - near_plane)
@@ -107,15 +93,3 @@ def interlevel_loss(weights_list, ray_samples_list, near_plane, far_plane):
         wp = weights[..., 0]  # (num_rays, num_samples)
         loss_interlevel += torch.mean(lossfun_outer(c, w, cp, wp))
     return loss_interlevel
-
-
-def distortion_loss(weights_list, ray_samples_list, near_plane, far_plane):
-    """Calculates the distortion loss in the MipNeRF-360 paper.
-
-    https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/model.py#L526
-    https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/train_utils.py#L147
-    """
-    c = ray_samples_to_sdist(ray_samples_list[-1], near_plane, far_plane).detach()
-    w = weights_list[-1][..., 0]
-    loss = torch.mean(lossfun_distortion(c, w))
-    return loss