| 522 | return {'pred_gs': pred_gs, 'gt_depth': target_depth, 'gt_depth_mask': mask}, prior_params |
| 523 | |
| 524 | class SplatPredictor(nn.Module): |
| 525 | def __init__(self, opt: Options, **model_kwargs): |
| 526 | super().__init__() |
| 527 | self.opt = opt |
| 528 | self.decoder = SplatDecoder(opt, **model_kwargs) |
| 529 | self.patch_size = opt.patch_size |
| 530 | |
| 531 | # # Upsampler |
| 532 | self.decoder_ratio = opt.decoder_ratio |
| 533 | if self.decoder_ratio > 0: |
| 534 | self.gaussian_upsampler = GaussianUpsampler(width=opt.hidden_dim, |
| 535 | ch_decay=2, |
| 536 | up_ratio=(2**self.decoder_ratio), |
| 537 | low_channels=opt.decoder_dim, |
| 538 | window_size=opt.window_size, |
| 539 | opt=opt) |
| 540 | transformer_dim = self.gaussian_upsampler.out_channels |
| 541 | else: |
| 542 | transformer_dim = opt.hidden_dim |
| 543 | |
| 544 | self.gs_predictor = GSPredictor(opt, **model_kwargs) |
| 545 | self.gs_predictor.predictor.keep_dynamic = True # keep dynamic for post combination |
| 546 | |
| 547 | if opt.use_dino: |
| 548 | self.condition_encoder = Dinov2Wrapper(model_name="dinov2_vitb14_reg") |
| 549 | self.condition_dim = 768 |
| 550 | |
| 551 | self.gs_dynamic_predictor = GSDynamicDecoder(opt, transformer_dim=self.gs_predictor.predictor.embed_dim, bias=True) |
| 552 | self.encoder_proj = nn.Linear(192, 192, bias=False) # hard code for now |
| 553 | |
| 554 | def state_dict(self, **kwargs): |
| 555 | # remove the condition encoder from the state dict |
| 556 | state_dict = super().state_dict(**kwargs) |
| 557 | for k in list(state_dict.keys()): |
| 558 | if "condition_encoder" in k: |
| 559 | del state_dict[k] |
| 560 | return state_dict |
| 561 | |
| 562 | def train(self, mode=True): |
| 563 | super().train(mode) |
| 564 | if self.opt.use_dino: |
| 565 | self.condition_encoder.eval() |
| 566 | self.gs_predictor.eval() |
| 567 | return self |
| 568 | |
| 569 | def _freeze_predictor(self): |
| 570 | self.gs_predictor.eval() |
| 571 | for name, param in self.gs_predictor.named_parameters(): |
| 572 | if "dynamic" in name: |
| 573 | param.requires_grad = True |
| 574 | elif "encoder" in name: |
| 575 | param.requires_grad = False |
| 576 | elif "upsampler" in name: |
| 577 | param.requires_grad = False |
| 578 | elif "gs_layer" in name: |
| 579 | param.requires_grad = False |
| 580 | else: |
| 581 | param.requires_grad = False |