| 218 | |
| 219 | # Model |
| 220 | class NeRFW(nn.Module): |
| 221 | def __init__(self, typ, |
| 222 | D=8, W=256, skips=[4], |
| 223 | in_channels_xyz=63, in_channels_dir=27, |
| 224 | encode_appearance=False, in_channels_a=48, |
| 225 | encode_transient=False, in_channels_t=16, |
| 226 | beta_min=0.1, out_ch_size=3): |
| 227 | """ |
| 228 | ---Parameters for the original NeRF--- |
| 229 | D: number of layers for density (sigma) encoder |
| 230 | W: number of hidden units in each layer |
| 231 | skips: add skip connection in the Dth layer |
| 232 | in_channels_xyz: number of input channels for xyz (3+3*10*2=63 by default) |
| 233 | in_channels_dir: number of input channels for direction (3+3*4*2=27 by default) |
| 234 | in_channels_t: number of input channels for t |
| 235 | |
| 236 | ---Parameters for NeRF-W (used in fine model only as per section 4.3)--- |
| 237 | ---cf. Figure 3 of the paper--- |
| 238 | encode_appearance: whether to add appearance encoding as input (NeRF-A) |
| 239 | in_channels_a: appearance embedding dimension. n^(a) in the paper |
| 240 | encode_transient: whether to add transient encoding as input (NeRF-U) |
| 241 | in_channels_t: transient embedding dimension. n^(tau) in the paper |
| 242 | beta_min: minimum pixel color variance |
| 243 | """ |
| 244 | super().__init__() |
| 245 | torch.manual_seed(0) |
| 246 | self.typ = typ |
| 247 | self.D = D |
| 248 | self.W = W |
| 249 | self.skips = skips |
| 250 | self.in_channels_xyz = in_channels_xyz |
| 251 | self.in_channels_dir = in_channels_dir |
| 252 | self.encode_appearance = False if typ=='coarse' else encode_appearance |
| 253 | self.in_channels_a = in_channels_a if encode_appearance else 0 |
| 254 | self.encode_transient = False if typ=='coarse' else encode_transient |
| 255 | self.in_channels_t = in_channels_t |
| 256 | self.beta_min = beta_min |
| 257 | |
| 258 | # xyz encoding layers |
| 259 | for i in range(D): |
| 260 | if i == 0: |
| 261 | layer = nn.Linear(in_channels_xyz, W) |
| 262 | elif i in skips: |
| 263 | layer = nn.Linear(W+in_channels_xyz, W) |
| 264 | else: |
| 265 | layer = nn.Linear(W, W) |
| 266 | layer = nn.Sequential(layer, nn.ReLU(True)) |
| 267 | setattr(self, f"xyz_encoding_{i+1}", layer) |
| 268 | self.xyz_encoding_final = nn.Linear(W, W) |
| 269 | |
| 270 | # direction encoding layers |
| 271 | self.dir_encoding = nn.Sequential( |
| 272 | nn.Linear(W+in_channels_dir+self.in_channels_a, W//2), nn.ReLU(True)) |
| 273 | |
| 274 | # static output layers |
| 275 | self.static_sigma = nn.Sequential(nn.Linear(W, 1), nn.Softplus()) |
| 276 | |
| 277 | if out_ch_size == 3: #NeRF, NeRFW output rgb |