forked from learning3d/assignment3
-
Notifications
You must be signed in to change notification settings - Fork 0
/
renderer.py
107 lines (84 loc) · 2.84 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
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 density
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
renderer_dict = {
'volume': VolumeRenderer
}