| 448 | return v |
| 449 | |
| 450 | class GSPredictor(nn.Module): |
| 451 | def __init__(self, opt: Options, **model_kwargs): |
| 452 | super(GSPredictor, self).__init__() |
| 453 | self.opt = opt |
| 454 | self.encoder = GSEncoder(opt, **model_kwargs) |
| 455 | self.patch_size = opt.patch_size |
| 456 | |
| 457 | if len(opt.down_resolution) > 0: |
| 458 | self.actual_input_res = opt.down_resolution |
| 459 | else: |
| 460 | self.actual_input_res = (opt.image_height, opt.image_width) |
| 461 | |
| 462 | # Upsampler |
| 463 | self.decoder_ratio = opt.decoder_ratio |
| 464 | if self.decoder_ratio > 0: |
| 465 | self.gaussian_upsampler = GaussianUpsampler(width=opt.hidden_dim, |
| 466 | ch_decay=2, |
| 467 | up_ratio=(2**self.decoder_ratio), |
| 468 | low_channels=opt.decoder_dim, |
| 469 | window_size=opt.window_size, |
| 470 | opt=opt) |
| 471 | transformer_dim = self.gaussian_upsampler.out_channels |
| 472 | else: |
| 473 | transformer_dim = opt.hidden_dim |
| 474 | # check if exist gaussian upsampler |
| 475 | # if hasattr(self.encoder, "gaussian_upsampler"): |
| 476 | # transformer_dim = self.encoder.gaussian_upsampler.out_channels |
| 477 | # else: |
| 478 | # transformer_dim = self.opt.hidden_dim |
| 479 | if self.opt.use_pm: |
| 480 | self.predictor = GSPMDecoder(opt, transformer_dim=transformer_dim) |
| 481 | else: |
| 482 | raise NotImplementedError("Currently only support probabilistic predictor") |
| 483 | |
| 484 | def _freeze(self): |
| 485 | for name, param in self.named_parameters(): |
| 486 | if "xyz_static" in name: |
| 487 | param.requires_grad = True |
| 488 | else: |
| 489 | param.requires_grad = False |
| 490 | |
| 491 | def forward_encoder(self, frames, depths, cond_times=None): |
| 492 | mask = torch.ones_like(depths) # dummy mask |
| 493 | max_depth = depths.flatten(1).max(dim=1)[0][:, None, None, None, None] |
| 494 | min_depth = depths.flatten(1).min(dim=1)[0][:, None, None, None, None] |
| 495 | target_depth = depths # without normalization |
| 496 | input_depth = (depths - min_depth) / (max_depth - min_depth) |
| 497 | |
| 498 | frames = torch.cat([frames, input_depth], dim=2) |
| 499 | encoder_output = self.encoder(frames, cond_times) |
| 500 | |
| 501 | return encoder_output, target_depth, mask |
| 502 | |
| 503 | def upsampling(self, output): |
| 504 | input_views = output.shape[1] |
| 505 | if self.decoder_ratio > 0: |
| 506 | output = output.reshape(output.shape[0], -1, output.shape[-1]) # [B, V, N, D] -> [B, V*N, D] |
| 507 | output = self.gaussian_upsampler(output) # [B, V*N, D] |