-
Notifications
You must be signed in to change notification settings - Fork 19
/
renderer.py
276 lines (222 loc) · 7.78 KB
/
renderer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
import torch
from typing import List, Optional, Tuple
from pytorch3d.renderer.cameras import CamerasBase
# Volume renderer which integrates color and density along rays
# according to the equations defined in [Mildenhall et al. 2020]
class VolumeRenderer(torch.nn.Module):
def __init__(
self,
cfg
):
super().__init__()
self._chunk_size = cfg.chunk_size
self._white_background = cfg.white_background if 'white_background' in cfg else False
def _compute_weights(
self,
deltas,
rays_density: torch.Tensor,
eps: float = 1e-10
):
# TODO (1.5): Compute transmittance using the equation described in the README
pass
# TODO (1.5): Compute weight used for rendering from transmittance and alpha
return weights
def _aggregate(
self,
weights: torch.Tensor,
rays_feature: torch.Tensor
):
# TODO (1.5): Aggregate (weighted sum of) features using weights
pass
return feature
def forward(
self,
sampler,
implicit_fn,
ray_bundle,
):
B = ray_bundle.shape[0]
# Process the chunks of rays.
chunk_outputs = []
for chunk_start in range(0, B, self._chunk_size):
cur_ray_bundle = ray_bundle[chunk_start:chunk_start+self._chunk_size]
# Sample points along the ray
cur_ray_bundle = sampler(cur_ray_bundle)
n_pts = cur_ray_bundle.sample_shape[1]
# Call implicit function with sample points
implicit_output = implicit_fn(cur_ray_bundle)
density = implicit_output['density']
feature = implicit_output['feature']
# Compute length of each ray segment
depth_values = cur_ray_bundle.sample_lengths[..., 0]
deltas = torch.cat(
(
depth_values[..., 1:] - depth_values[..., :-1],
1e10 * torch.ones_like(depth_values[..., :1]),
),
dim=-1,
)[..., None]
# Compute aggregation weights
weights = self._compute_weights(
deltas.view(-1, n_pts, 1),
density.view(-1, n_pts, 1)
)
# TODO (1.5): Render (color) features using weights
pass
# TODO (1.5): Render depth map
pass
# Return
cur_out = {
'feature': feature,
'depth': depth,
}
chunk_outputs.append(cur_out)
# Concatenate chunk outputs
out = {
k: torch.cat(
[chunk_out[k] for chunk_out in chunk_outputs],
dim=0
) for k in chunk_outputs[0].keys()
}
return out
# Volume renderer which integrates color and density along rays
# according to the equations defined in [Mildenhall et al. 2020]
class SphereTracingRenderer(torch.nn.Module):
def __init__(
self,
cfg
):
super().__init__()
self._chunk_size = cfg.chunk_size
self.near = cfg.near
self.far = cfg.far
self.max_iters = cfg.max_iters
def sphere_tracing(
self,
implicit_fn,
origins, # Nx3
directions, # Nx3
):
'''
Input:
implicit_fn: a module that computes a SDF at a query point
origins: N_rays X 3
directions: N_rays X 3
Output:
points: N_rays X 3 points indicating ray-surface intersections. For rays that do not intersect the surface,
the point can be arbitrary.
mask: N_rays X 1 (boolean tensor) denoting which of the input rays intersect the surface.
'''
# TODO (Q5): Implement sphere tracing
# 1) Iteratively update points and distance to the closest surface
# in order to compute intersection points of rays with the implicit surface
# 2) Maintain a mask with the same batch dimension as the ray origins,
# indicating which points hit the surface, and which do not
pass
def forward(
self,
sampler,
implicit_fn,
ray_bundle,
light_dir=None
):
B = ray_bundle.shape[0]
# Process the chunks of rays.
chunk_outputs = []
for chunk_start in range(0, B, self._chunk_size):
cur_ray_bundle = ray_bundle[chunk_start:chunk_start+self._chunk_size]
points, mask = self.sphere_tracing(
implicit_fn,
cur_ray_bundle.origins,
cur_ray_bundle.directions
)
mask = mask.repeat(1,3)
isect_points = points[mask].view(-1, 3)
# Get color from implicit function with intersection points
isect_color = implicit_fn.get_color(isect_points)
# Return
color = torch.zeros_like(cur_ray_bundle.origins)
color[mask] = isect_color.view(-1)
cur_out = {
'color': color.view(-1, 3),
}
chunk_outputs.append(cur_out)
# Concatenate chunk outputs
out = {
k: torch.cat(
[chunk_out[k] for chunk_out in chunk_outputs],
dim=0
) for k in chunk_outputs[0].keys()
}
return out
def sdf_to_density(signed_distance, alpha, beta):
# TODO (Q7): Convert signed distance to density with alpha, beta parameters
pass
class VolumeSDFRenderer(VolumeRenderer):
def __init__(
self,
cfg
):
super().__init__(cfg)
self._chunk_size = cfg.chunk_size
self._white_background = cfg.white_background if 'white_background' in cfg else False
self.alpha = cfg.alpha
self.beta = cfg.beta
self.cfg = cfg
def forward(
self,
sampler,
implicit_fn,
ray_bundle,
light_dir=None
):
B = ray_bundle.shape[0]
# Process the chunks of rays.
chunk_outputs = []
for chunk_start in range(0, B, self._chunk_size):
cur_ray_bundle = ray_bundle[chunk_start:chunk_start+self._chunk_size]
# Sample points along the ray
cur_ray_bundle = sampler(cur_ray_bundle)
n_pts = cur_ray_bundle.sample_shape[1]
# Call implicit function with sample points
distance, color = implicit_fn.get_distance_color(cur_ray_bundle.sample_points)
density = None # TODO (Q7): convert SDF to density
# Compute length of each ray segment
depth_values = cur_ray_bundle.sample_lengths[..., 0]
deltas = torch.cat(
(
depth_values[..., 1:] - depth_values[..., :-1],
1e10 * torch.ones_like(depth_values[..., :1]),
),
dim=-1,
)[..., None]
# Compute aggregation weights
weights = self._compute_weights(
deltas.view(-1, n_pts, 1),
density.view(-1, n_pts, 1)
)
geometry_color = torch.zeros_like(color)
# Compute color
color = self._aggregate(
weights,
color.view(-1, n_pts, color.shape[-1])
)
# Return
cur_out = {
'color': color,
"geometry": geometry_color
}
chunk_outputs.append(cur_out)
# Concatenate chunk outputs
out = {
k: torch.cat(
[chunk_out[k] for chunk_out in chunk_outputs],
dim=0
) for k in chunk_outputs[0].keys()
}
return out
renderer_dict = {
'volume': VolumeRenderer,
'sphere_tracing': SphereTracingRenderer,
'volume_sdf': VolumeSDFRenderer
}