From 9caf9934b684d0b54826399230bfb356592a3a33 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Mon, 26 Sep 2022 15:13:17 -0700 Subject: [PATCH] use our distortion loss (#543) swap out with our distortion loss --- nerfactory/models/modules/ray_losses.py | 21 +++++++++++++++----- nerfactory/models/proposal.py | 18 +++++++++-------- nerfactory/optimizers/loss.py | 26 ------------------------- 3 files changed, 26 insertions(+), 39 deletions(-) diff --git a/nerfactory/models/modules/ray_losses.py b/nerfactory/models/modules/ray_losses.py index 9724f6112..f74937093 100644 --- a/nerfactory/models/modules/ray_losses.py +++ b/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 diff --git a/nerfactory/models/proposal.py b/nerfactory/models/proposal.py index a99314ffa..875f7e9e1 100644 --- a/nerfactory/models/proposal.py +++ b/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} diff --git a/nerfactory/optimizers/loss.py b/nerfactory/optimizers/loss.py index 6f79be8f8..57b9a7633 100644 --- a/nerfactory/optimizers/loss.py +++ b/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