From 611aba9a202886cf493d841610ae21c27401b923 Mon Sep 17 00:00:00 2001
From: Jeremy Reizenstein <reizenstein@meta.com>
Date: Sun, 23 Oct 2022 02:47:09 -0700
Subject: [PATCH] replicate_last_interval in raymarcher

Summary: Add option to flat pad the last delta. Might to help when training on rgb only.

Reviewed By: shapovalov

Differential Revision: D40587475

fbshipit-source-id: c763fa38948600ea532c730538dc4ff29d2c3e0a
---
 .../implicitron_trainer/tests/experiment.yaml |  2 ++
 .../implicitron/models/renderer/raymarcher.py | 23 +++++++++++--------
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/projects/implicitron_trainer/tests/experiment.yaml b/projects/implicitron_trainer/tests/experiment.yaml
index f9e6e329..87236c93 100644
--- a/projects/implicitron_trainer/tests/experiment.yaml
+++ b/projects/implicitron_trainer/tests/experiment.yaml
@@ -235,6 +235,7 @@ model_factory_ImplicitronModelFactory_args:
         surface_thickness: 1
         bg_color:
         - 0.0
+        replicate_last_interval: False
         background_opacity: 0.0
         density_relu: true
         blend_output: false
@@ -242,6 +243,7 @@ model_factory_ImplicitronModelFactory_args:
         surface_thickness: 1
         bg_color:
         - 0.0
+        replicate_last_interval: False
         background_opacity: 10000000000.0
         density_relu: true
         blend_output: false
diff --git a/pytorch3d/implicitron/models/renderer/raymarcher.py b/pytorch3d/implicitron/models/renderer/raymarcher.py
index b9cd2980..37ddd6c5 100644
--- a/pytorch3d/implicitron/models/renderer/raymarcher.py
+++ b/pytorch3d/implicitron/models/renderer/raymarcher.py
@@ -55,8 +55,11 @@ class AccumulativeRaymarcherBase(RaymarcherBase, torch.nn.Module):
         surface_thickness: The thickness of the raymarched surface.
         bg_color: The background color. A tuple of either 1 element or of D elements,
             where D matches the feature dimensionality; it is broadcast when necessary.
-        background_opacity: The raw opacity value (i.e. before exponentiation)
-            of the background.
+        replicate_last_interval: If True, the ray length assigned to the last interval
+            for the opacity delta calculation is copied from the penultimate interval.
+        background_opacity: The length over which the last raw opacity value
+            (i.e. before exponentiation) is considered to apply, for the delta
+            calculation. Ignored if replicate_last_interval=True.
         density_relu: If `True`, passes the input density through ReLU before
             raymarching.
         blend_output: If `True`, alpha-blends the output renders with the
@@ -76,6 +79,7 @@ class AccumulativeRaymarcherBase(RaymarcherBase, torch.nn.Module):
 
     surface_thickness: int = 1
     bg_color: Tuple[float, ...] = (0.0,)
+    replicate_last_interval: bool = False
     background_opacity: float = 0.0
     density_relu: bool = True
     blend_output: bool = False
@@ -151,13 +155,14 @@ class AccumulativeRaymarcherBase(RaymarcherBase, torch.nn.Module):
             density_1d=True,
         )
 
-        deltas = torch.cat(
-            (
-                ray_lengths[..., 1:] - ray_lengths[..., :-1],
-                self.background_opacity * torch.ones_like(ray_lengths[..., :1]),
-            ),
-            dim=-1,
-        )
+        ray_lengths_diffs = ray_lengths[..., 1:] - ray_lengths[..., :-1]
+        if self.replicate_last_interval:
+            last_interval = ray_lengths_diffs[..., -1:]
+        else:
+            last_interval = torch.full_like(
+                ray_lengths[..., :1], self.background_opacity
+            )
+        deltas = torch.cat((ray_lengths_diffs, last_interval), dim=-1)
 
         rays_densities = rays_densities[..., 0]