| 232 | layer_dict[f"{key}_scale"] = pred_scale |
| 233 | |
| 234 | def key_activation(self, v: torch.Tensor, key=''): |
| 235 | # [B, N, D] |
| 236 | B = v.shape[0] |
| 237 | assert v.isnan().sum() == 0, f"NaN detected in {key}" |
| 238 | v = v.type(torch.float32) |
| 239 | if key == "xyz_static": |
| 240 | # (B, N, 3) v shape |
| 241 | v = torch.tanh(v) |
| 242 | x_pred, y_pred, z_pred = v.chunk(3, dim=-1) |
| 243 | y_val = 0.5 + y_pred * 0.5 |
| 244 | # predict x and z with pixel position |
| 245 | x_offset = self.x_max * x_pred |
| 246 | z_offset = self.z_max * z_pred |
| 247 | x_map = self.x_map.repeat(self.opt.input_frames).reshape(1, -1, 1) |
| 248 | z_map = self.z_map.repeat(self.opt.input_frames).reshape(1, -1, 1) |
| 249 | x_val = x_map + x_offset |
| 250 | z_val = z_map + z_offset |
| 251 | v = torch.cat([x_val, y_val, z_val], dim=-1) |
| 252 | elif key == "scale": |
| 253 | v = 0.1 * F.softplus(v) |
| 254 | elif key == "rot_static": |
| 255 | v = F.normalize(v, dim=-1) |
| 256 | elif key == "opacity": |
| 257 | if self.opacity_activation == "sigmoid": |
| 258 | v = torch.sigmoid(v) |
| 259 | v = 0.05 + 0.95 * v |
| 260 | else: |
| 261 | v = F.relu(v + 10) |
| 262 | elif key == "shs": |
| 263 | pass |
| 264 | elif key == "rgb": |
| 265 | v = torch.sigmoid(v) |
| 266 | else: |
| 267 | raise NotImplementedError |
| 268 | if v.dim() == 3: |
| 269 | v = v.repeat(1, 1, self.n_sample).reshape(v.shape[0], -1, v.shape[-1]) # maybe incorrect |
| 270 | elif v.dim() == 4: |
| 271 | v = v.repeat(1, 1, 1, self.n_sample).reshape(v.shape[0], -1, v.shape[-1]) |
| 272 | else: |
| 273 | raise NotImplementedError |
| 274 | return v |
| 275 | |
| 276 | @autocast('cuda', enabled=False) |
| 277 | def forward(self, feats, timestamp=None): |