diff --git a/configs/graph_mipnerf_360.yaml b/configs/graph_mipnerf_360.yaml index d5fdbaa7a..3817d540a 100644 --- a/configs/graph_mipnerf_360.yaml +++ b/configs/graph_mipnerf_360.yaml @@ -27,4 +27,7 @@ graph: collider_config: _target_: pyrad.graphs.modules.scene_colliders.NearFarCollider near_plane: 0.5 - far_plane: 20.0 \ No newline at end of file + far_plane: 20.0 + loss_coefficients: + ray_loss_coarse: 1.0 + ray_loss_fine: 1.0 \ No newline at end of file diff --git a/docs/reference/graphs/modules/index.rst b/docs/reference/graphs/modules/index.rst index cfeea81a5..1e401b782 100644 --- a/docs/reference/graphs/modules/index.rst +++ b/docs/reference/graphs/modules/index.rst @@ -7,4 +7,5 @@ TODO: High level description of graph modules and how they are used. .. toctree:: - ray_sampler \ No newline at end of file + ray_sampler + ray_losses \ No newline at end of file diff --git a/docs/reference/graphs/modules/ray_losses.rst b/docs/reference/graphs/modules/ray_losses.rst new file mode 100644 index 000000000..3aa1de992 --- /dev/null +++ b/docs/reference/graphs/modules/ray_losses.rst @@ -0,0 +1,8 @@ +.. _ray_losses: + +Ray Losses +=================== + +.. automodule:: pyrad.graphs.modules.ray_losses + :members: + :show-inheritance: \ No newline at end of file diff --git a/pyrad/graphs/mipnerf_360.py b/pyrad/graphs/mipnerf_360.py index 147094309..c157fe98c 100644 --- a/pyrad/graphs/mipnerf_360.py +++ b/pyrad/graphs/mipnerf_360.py @@ -30,6 +30,7 @@ from pyrad.fields.modules.field_heads import FieldHeadNames from pyrad.fields.modules.spatial_distortions import SceneContraction from pyrad.fields.nerf_field import NeRFField from pyrad.graphs.base import Graph +from pyrad.graphs.modules.ray_losses import interval_loss from pyrad.optimizers.loss import MSELoss from pyrad.graphs.modules.ray_sampler import PDFSampler, UniformSampler from pyrad.renderers.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer @@ -111,6 +112,7 @@ class MipNerf360Graph(Graph): ) accumulation_coarse = self.renderer_accumulation(weights_coarse) depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform) + ray_loss_coarse = interval_loss(ray_samples_uniform, field_outputs_coarse[FieldHeadNames.DENSITY]) # pdf sampling ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) @@ -124,6 +126,7 @@ class MipNerf360Graph(Graph): ) accumulation_fine = self.renderer_accumulation(weights_fine) depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf) + ray_loss_fine = interval_loss(ray_samples_pdf, field_outputs_fine[FieldHeadNames.DENSITY]) outputs = { "rgb_coarse": rgb_coarse, @@ -132,6 +135,8 @@ class MipNerf360Graph(Graph): "accumulation_fine": accumulation_fine, "depth_coarse": depth_coarse, "depth_fine": depth_fine, + "ray_loss_coarse": ray_loss_coarse, + "ray_loss_fine": ray_loss_fine, } return outputs @@ -139,7 +144,12 @@ class MipNerf360Graph(Graph): image = batch["image"] rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"]) rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"]) - loss_dict = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine} + loss_dict = { + "rgb_loss_coarse": rgb_loss_coarse, + "rgb_loss_fine": rgb_loss_fine, + "ray_loss_coarse": torch.mean(outputs["ray_loss_coarse"]), + "ray_loss_fine": torch.mean(outputs["ray_loss_fine"]), + } return loss_dict def log_test_image_outputs(self, image_idx, step, batch, outputs): @@ -183,6 +193,12 @@ class MipNerf360Graph(Graph): writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(fine_psnr), step=step) writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(fine_ssim), step=step) writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(fine_lpips), step=step) + writer.put_scalar( + name=f"ray_loss_coarse/val_{image_idx}", scalar=float(torch.mean(outputs["ray_loss_coarse"])), step=step + ) + writer.put_scalar( + name=f"ray_loss_fine/val_{image_idx}", scalar=float(torch.mean(outputs["ray_loss_fine"])), step=step + ) writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step) diff --git a/pyrad/graphs/modules/ray_losses.py b/pyrad/graphs/modules/ray_losses.py new file mode 100644 index 000000000..51dbaae20 --- /dev/null +++ b/pyrad/graphs/modules/ray_losses.py @@ -0,0 +1,53 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Collection of Ray Losses""" + +import torch +from torchtyping import TensorType + +from pyrad.cameras.rays import RaySamples + + +def interval_loss(ray_samples: RaySamples, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., 1]: + """Interval loss origially proposed in MipNeRF-360. + + .. math:: + + \\mathcal{L}(\\mathbf{s}, \\mathbf{w}) =\\iint\\limits_{-\\infty}^{\\,\\,\\,\\infty} + \\mathbf{w}_\\mathbf{s}(u)\\mathbf{w}_\\mathbf{s}(v)|u - v|\\,d_{u}\\,d_{v} + + where :math:`\\mathbf{w}_\\mathbf{s}(u)=\\sum_i w_i \\mathbb{1}_{[\\mathbf{s}_i, \\mathbf{s}_{i+1})}(u)` + is the weight at location :math:`u` between bin locations :math:`s_i` and :math:`s_{i+1}`. + + Args: + ray_samples (RaySamples): Ray samples to compute loss over + densities (TensorType[..., "num_samples", 1]): Predicted sample densities + + Returns: + TensorType[..., 1]: Interval Loss. + """ + + # Compute the weight at each sample location + weights = ray_samples.get_weights(densities) + + midpoints = (ray_samples.bin_starts + ray_samples.bin_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.bin_ends - ray_samples.bin_starts), dim=-2) + + return loss