| 39 | |
| 40 | |
| 41 | class Model(nn.Module): |
| 42 | def __init__( |
| 43 | self, |
| 44 | cfg: ModelCfg, |
| 45 | num_frames: int | None = None, |
| 46 | image_shape: tuple[int, int] | None = None, |
| 47 | ) -> None: |
| 48 | super().__init__() |
| 49 | self.cfg = cfg |
| 50 | self.backbone = get_backbone(cfg.backbone, num_frames, image_shape) |
| 51 | self.intrinsics = get_intrinsics(cfg.intrinsics) |
| 52 | self.extrinsics = get_extrinsics(cfg.extrinsics, num_frames) |
| 53 | |
| 54 | def forward( |
| 55 | self, |
| 56 | batch: Batch, |
| 57 | flows: Flows, |
| 58 | global_step: int, |
| 59 | ) -> ModelOutput: |
| 60 | device = batch.videos.device |
| 61 | _, _, _, h, w = batch.videos.shape |
| 62 | |
| 63 | # Run the backbone, which provides depths and correspondence weights. |
| 64 | backbone_out = self.backbone.forward(batch, flows) |
| 65 | |
| 66 | # Allow the correspondence weights to be ignored as an ablation. |
| 67 | if not self.cfg.use_correspondence_weights: |
| 68 | backbone_out.weights = torch.ones_like(backbone_out.weights) |
| 69 | |
| 70 | # Compute the intrinsics. |
| 71 | intrinsics = self.intrinsics.forward(batch, flows, backbone_out, global_step) |
| 72 | |
| 73 | # Use the intrinsics to calculate camera-space surfaces (point clouds). |
| 74 | xy, _ = sample_image_grid((h, w), device=device) |
| 75 | surfaces = unproject( |
| 76 | xy, |
| 77 | backbone_out.depths, |
| 78 | rearrange(intrinsics, "b f i j -> b f () () i j"), |
| 79 | ) |
| 80 | |
| 81 | # Finally, compute the extrinsics. |
| 82 | extrinsics = self.extrinsics.forward(batch, flows, backbone_out, surfaces) |
| 83 | |
| 84 | return ModelOutput( |
| 85 | backbone_out.depths, |
| 86 | surfaces, |
| 87 | intrinsics, |
| 88 | extrinsics, |
| 89 | backbone_out.weights, |
| 90 | ) |
| 91 | |
| 92 | @torch.no_grad() |
| 93 | def export( |
| 94 | self, |
| 95 | batch: Batch, |
| 96 | flows: Flows, |
| 97 | global_step: int, |
| 98 | ) -> ModelExports: |