Class MonodepthWithEncodingAdaptor

src/sharp/models/monodepth.py:167–251

Monodepth model with feature maps.

class MonodepthWithEncodingAdaptor(nn.Module):
    """Monodepth model with feature maps."""

    def __init__(
        self,
        monodepth_predictor: MonodepthDensePredictionTransformer,
        return_encoder_features: bool,
        return_decoder_features: bool,
        num_monodepth_layers: int,
        sorting_monodepth: bool,
    ):
        """Initialize MonodepthWithEncodingAdaptor.

        Args:
            monodepth_predictor: The monodepth model.
            return_encoder_features: Whether to return encoder features from monodepth model.
            return_decoder_features: Whether to return decoder features from monodepth model.
            num_monodepth_layers: How many layers the monodepth model predicts.
            sorting_monodepth: Whether to sort the monodepth output (for two layer monodepth).
        """
        super().__init__()
        self.monodepth_predictor = monodepth_predictor
        self.return_encoder_features = return_encoder_features
        self.return_decoder_features = return_decoder_features
        self.num_monodepth_layers = num_monodepth_layers
        self.sorting_monodepth = sorting_monodepth

    def forward(self, image: torch.Tensor) -> MonodepthOutput:
        """Process image and return disparity and feature maps."""
        inputs = self.monodepth_predictor.normalizer(image)
        encoder_output = self.monodepth_predictor.encoder(inputs)

        num_encoder_features = len(self.monodepth_predictor.encoder.dims_encoder)

        # NOTE: whether intermediate features are empty have already been decided
        # in monodepth_predictor during create_monodepth_dpt.
        encoder_features = encoder_output[:num_encoder_features]
        intermediate_features = encoder_output[num_encoder_features:]
        decoder_features = self.monodepth_predictor.decoder(encoder_features)
        disparity = self.monodepth_predictor.head(decoder_features)

        # We cannot use disparity.shape[1], otherwise the tracer will fail.
        if self.num_monodepth_layers == 2 and self.sorting_monodepth:
            first_layer_disparity = disparity.max(dim=1, keepdims=True).values
            second_layer_disparity = disparity.min(dim=1, keepdims=True).values
            disparity = torch.cat([first_layer_disparity, second_layer_disparity], dim=1)

        output_features = []
        if self.return_encoder_features:
            output_features.extend(encoder_features)

        if self.return_decoder_features:
            output_features.append(decoder_features)

        return MonodepthOutput(
            disparity=disparity,
            encoder_features=encoder_features,
            decoder_features=decoder_features,