google-research

1
--- external/gsn/models/nerf_utils.py	2023-03-28 16:59:06.690064402 +0000
2
+++ external_reference/gsn/models/nerf_utils.py	2023-03-28 16:00:21.946531010 +0000
3
@@ -141,7 +141,7 @@
4
 
5
     return pts, viewdirs, z_vals, rd, ro
6
 
7
-
8
+# CHANGED: modified to use inverse-depth (disparity), and change thickness of final point on ray
9
 def volume_render_radiance_field(
10
     rgb,
11
     occupancy,
12
@@ -151,16 +151,13 @@
13
     alpha_activation='relu',
14
     activate_rgb=True,
15
     density_bias=0,
16
+    use_disp=True,
17
 ):
18
 
19
-    one_e_10 = torch.tensor([1e10], dtype=ray_directions.dtype, device=ray_directions.device)
20
-    dists = torch.cat(
21
-        (
22
-            depth_values[..., 1:] - depth_values[..., :-1],
23
-            one_e_10.expand(depth_values[..., :1].shape),
24
-        ),
25
-        dim=-1,
26
-    )
27
+    # CHANGED: assume that the last sample is the same thickness as previous
28
+    # samples, rather than one_e_10 (which makes it hard for model to learn transparency)
29
+    dists = depth_values[..., 1:] - depth_values[..., :-1]
30
+    dists = torch.cat([dists, dists[..., [-1]]], dim=-1)
31
     dists = dists * ray_directions[..., None, :].norm(p=2, dim=-1)
32
 
33
     noise = 0.0
34
@@ -184,7 +181,9 @@
35
     weights = alpha * cumprod_exclusive(1.0 - alpha + 1e-10)
36
 
37
     if activate_rgb:
38
-        rgb = torch.sigmoid(rgb)
39
+        # CHANGED: use [-1, 1] activation
40
+        rgb = torch.tanh(rgb/2) # scaled sigmoid [-1, 1]
41
+        # rgb = torch.sigmoid(rgb)
42
 
43
         # widened sigmoid from https://github.com/google/mipnerf/blob/main/internal/models.py#L123
44
         rgb_padding = 0.001
45
@@ -193,7 +192,11 @@
46
     rgb_map = weights[..., None] * rgb
47
     rgb_map = rgb_map.sum(dim=-2)
48
 
49
-    depth_map = weights * depth_values
50
+    # CHANGED: return disparity rather than depth
51
+    if use_disp:
52
+        depth_map = weights * 1/depth_values
53
+    else:
54
+        depth_map = weights * depth_values
55
     depth_map = depth_map.sum(dim=-1)
56
 
57
     acc_map = weights.sum(dim=-1)
58
@@ -205,7 +208,13 @@
59
         torch.log(0.1 + alpha.view(alpha.size(0), -1)) + torch.log(0.1 + 1.0 - alpha.view(alpha.size(0), -1)) - -2.20727
60
     )
61
 
62
-    return rgb_map, disp_map, acc_map, weights, depth_map, occupancy_prior
63
+    # used for opacity regularization
64
+    extras = {
65
+        'dists': dists,
66
+        'alpha': alpha,
67
+    }
68
+
69
+    return rgb_map, disp_map, acc_map, weights, depth_map, occupancy_prior, extras
70
 
71
 
72
 def sample_pdf_2(bins, weights, num_samples, det=False):
73