| 70 | |
| 71 | # Model |
| 72 | class NeRF(nn.Module): |
| 73 | def __init__(self, D=8, W=256, input_ch=3, input_ch_views=3, output_ch=4, skips=[4], use_viewdirs=False, beta_min=0.0): |
| 74 | """ |
| 75 | """ |
| 76 | super(NeRF, self).__init__() |
| 77 | self.D = D |
| 78 | self.W = W |
| 79 | self.input_ch = input_ch |
| 80 | self.input_ch_views = input_ch_views |
| 81 | self.skips = skips |
| 82 | self.use_viewdirs = use_viewdirs |
| 83 | self.beta_min = beta_min |
| 84 | self.pts_linears = nn.ModuleList( |
| 85 | [nn.Linear(input_ch, W)] + [nn.Linear(W, W) if i not in self.skips else nn.Linear(W + input_ch, W) for i in range(D-1)]) |
| 86 | |
| 87 | ### Implementation according to the official code release (https://github.com/bmild/nerf/blob/master/run_nerf_helpers.py#L104-L105) |
| 88 | self.views_linears = nn.ModuleList([nn.Linear(input_ch_views + W, W//2)]) |
| 89 | |
| 90 | ### Implementation according to the paper |
| 91 | # self.views_linears = nn.ModuleList( |
| 92 | # [nn.Linear(input_ch_views + W, W//2)] + [nn.Linear(W//2, W//2) for i in range(D//2)]) |
| 93 | |
| 94 | if use_viewdirs: |
| 95 | self.feature_linear = nn.Linear(W, W) |
| 96 | self.alpha_linear = nn.Linear(W, 1) |
| 97 | self.uncertainty_linear = nn.Linear(W, 1) |
| 98 | # self.act_uncertainty = nn.ReLU() |
| 99 | self.act_uncertainty = nn.Softplus() |
| 100 | self.rgb_linear = nn.Linear(W//2, 3) |
| 101 | else: |
| 102 | self.output_linear = nn.Linear(W, output_ch) |
| 103 | |
| 104 | def forward(self, x): |
| 105 | input_pts, input_views = torch.split(x, [self.input_ch, self.input_ch_views], dim=-1) |
| 106 | h = input_pts |
| 107 | for i, l in enumerate(self.pts_linears): |
| 108 | h = self.pts_linears[i](h) |
| 109 | h = F.relu(h) |
| 110 | if i in self.skips: |
| 111 | h = torch.cat([input_pts, h], -1) |
| 112 | |
| 113 | if self.use_viewdirs: |
| 114 | alpha = self.alpha_linear(h) |
| 115 | uncert = self.act_uncertainty(self.uncertainty_linear(h)) + self.beta_min |
| 116 | feature = self.feature_linear(h) |
| 117 | h = torch.cat([feature, input_views], -1) |
| 118 | |
| 119 | for i, l in enumerate(self.views_linears): |
| 120 | h = self.views_linears[i](h) |
| 121 | h = F.relu(h) |
| 122 | |
| 123 | rgb = self.rgb_linear(h) |
| 124 | outputs = torch.cat([rgb, alpha, uncert], -1) |
| 125 | else: |
| 126 | outputs = self.output_linear(h) |
| 127 | |
| 128 | return outputs |
| 129 | |