Class BackboneMidas

flowmap/model/backbone/backbone_midas.py:42–127

class BackboneMidas(Backbone[BackboneMidasCfg]):
    def __init__(
        self,
        cfg: BackboneMidasCfg,
        num_frames: int | None,
        image_shape: tuple[int, int] | None,
    ) -> None:
        super().__init__(cfg, num_frames=num_frames, image_shape=image_shape)
        self.midas = torch.hub.load(
            "intel-isl/MiDaS",
            cfg.model,
            pretrained=cfg.pretrained,
        )
        self.midas_out = self.midas.scratch.output_conv
        self.midas.scratch.output_conv = nn.Identity()

        # If a weight sensitivity is specified, don't learn weights.
        if cfg.weight_sensitivity is None:
            weight_channels = {
                "DPT_Large": 256,
                "MiDaS_small": 64,
            }[cfg.model]
            self.corr_weighter_perpoint = make_net([weight_channels * 2, 128, 64, 1])
        else:
            weights = torch.full((num_frames - 1, *image_shape), 0, dtype=torch.float32)
            self.weights = nn.Parameter(weights)

        if cfg.mapping == "exp":
            self.midas_out = nn.Sequential(*self.midas_out[:-2])

    def forward(self, batch: Batch, flows: Flows) -> BackboneOutput:
        device = batch.videos.device
        b, f, _, h, w = batch.videos.shape

        videos = rearrange(batch.videos, "b f c h w -> (b f) c h w")
        features = self.midas(videos)

        # This matches Cameron's original implementation.
        match self.cfg.mapping:
            case "original":
                depths = 1e3 / (self.midas_out(features) + 0.1)
            case "exp":
                depths = (self.midas_out(features) / 1000).exp() + 0.01

        features = F.interpolate(features, (h, w), mode="bilinear") / 20

        depths = rearrange(depths, "(b f) () h w -> b f h w", b=b, f=f)
        features = rearrange(features, "(b f) c h w -> b f c h w", b=b, f=f)

        # Compute correspondence weights.
        if self.cfg.weight_sensitivity is None:
            xy, _ = sample_image_grid((h, w), device)
            backward_weights = self.compute_correspondence_weights(
                self.grid_sample_features(earlier(features), xy + flows.backward),
                later(features),
            )
        else:
            backward_weights = (self.cfg.weight_sensitivity * self.weights).sigmoid()