(self, batch: Batch, flows: Flows)
| 70 | self.midas_out = nn.Sequential(*self.midas_out[:-2]) |
| 71 | |
| 72 | def forward(self, batch: Batch, flows: Flows) -> BackboneOutput: |
| 73 | device = batch.videos.device |
| 74 | b, f, _, h, w = batch.videos.shape |
| 75 | |
| 76 | videos = rearrange(batch.videos, "b f c h w -> (b f) c h w") |
| 77 | features = self.midas(videos) |
| 78 | |
| 79 | # This matches Cameron's original implementation. |
| 80 | match self.cfg.mapping: |
| 81 | case "original": |
| 82 | depths = 1e3 / (self.midas_out(features) + 0.1) |
| 83 | case "exp": |
| 84 | depths = (self.midas_out(features) / 1000).exp() + 0.01 |
| 85 | |
| 86 | features = F.interpolate(features, (h, w), mode="bilinear") / 20 |
| 87 | |
| 88 | depths = rearrange(depths, "(b f) () h w -> b f h w", b=b, f=f) |
| 89 | features = rearrange(features, "(b f) c h w -> b f c h w", b=b, f=f) |
| 90 | |
| 91 | # Compute correspondence weights. |
| 92 | if self.cfg.weight_sensitivity is None: |
| 93 | xy, _ = sample_image_grid((h, w), device) |
| 94 | backward_weights = self.compute_correspondence_weights( |
| 95 | self.grid_sample_features(earlier(features), xy + flows.backward), |
| 96 | later(features), |
| 97 | ) |
| 98 | else: |
| 99 | backward_weights = (self.cfg.weight_sensitivity * self.weights).sigmoid() |
| 100 | backward_weights = backward_weights[None] |
| 101 | |
| 102 | return BackboneOutput(depths, backward_weights) |
| 103 | |
| 104 | def compute_correspondence_weights( |
| 105 | self, |
nothing calls this directly
no test coverage detected