124 lines
4.2 KiB
Python
124 lines
4.2 KiB
Python
"""Contains Dense Transformer Prediction architecture.
|
|
|
|
Implements a variant of Vision Transformers for Dense Prediction, https://arxiv.org/abs/2103.13413
|
|
|
|
For licensing see accompanying LICENSE file.
|
|
Copyright (C) 2025 Apple Inc. All Rights Reserved.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from sharp.models.presets import (
|
|
MONODEPTH_ENCODER_DIMS_MAP,
|
|
MONODEPTH_HOOK_IDS_MAP,
|
|
ViTPreset,
|
|
)
|
|
|
|
from .base_encoder import BaseEncoder
|
|
from .spn_encoder import SlidingPyramidNetwork
|
|
from .vit_encoder import create_vit
|
|
|
|
|
|
def create_monodepth_encoder(
|
|
patch_encoder_preset: ViTPreset,
|
|
image_encoder_preset: ViTPreset,
|
|
use_patch_overlap: bool = True,
|
|
last_encoder: int = 256,
|
|
) -> SlidingPyramidNetwork:
|
|
"""Creates DepthDensePredictionTransformer model.
|
|
|
|
Args:
|
|
patch_encoder_preset: The preset patch encoder architecture in SPN.
|
|
image_encoder_preset: The preset image encoder architecture in SPN.
|
|
use_patch_overlap: Whether to use overlap between patches in SPN.
|
|
last_encoder: last number of encoder features.
|
|
"""
|
|
dims_encoder = [last_encoder] + MONODEPTH_ENCODER_DIMS_MAP[patch_encoder_preset]
|
|
patch_encoder_block_ids = MONODEPTH_HOOK_IDS_MAP[patch_encoder_preset]
|
|
|
|
patch_encoder = create_vit(
|
|
preset=patch_encoder_preset,
|
|
intermediate_features_ids=patch_encoder_block_ids,
|
|
# We always need to output intermediate features for assembly.
|
|
)
|
|
image_encoder = create_vit(
|
|
preset=image_encoder_preset,
|
|
intermediate_features_ids=None,
|
|
)
|
|
|
|
encoder = SlidingPyramidNetwork(
|
|
dims_encoder=dims_encoder,
|
|
patch_encoder=patch_encoder,
|
|
image_encoder=image_encoder,
|
|
use_patch_overlap=use_patch_overlap,
|
|
)
|
|
|
|
return encoder
|
|
|
|
|
|
class ProjectionModule(nn.Module):
|
|
"""Apply projection of features."""
|
|
|
|
def __init__(self, dims_in: list[int], dims_out: list[int]) -> None:
|
|
"""Initialize projection module."""
|
|
super().__init__()
|
|
if len(dims_in) != len(dims_out):
|
|
raise ValueError("Length of dims_in must be same as length of dims_out.")
|
|
self.convs = nn.ModuleList(
|
|
[nn.Conv2d(dim_in, dim_out, 1) for dim_in, dim_out in zip(dims_in, dims_out)]
|
|
)
|
|
|
|
def forward(self, encodings: list[torch.Tensor]) -> list[torch.Tensor]:
|
|
"""Apply projection module."""
|
|
if len(encodings) != len(self.convs):
|
|
raise ValueError("Number of encodings must be equal to number of projections.")
|
|
return [conv(encoding) for conv, encoding in zip(self.convs, encodings)]
|
|
|
|
|
|
class MonodepthFeatureEncoder(BaseEncoder):
|
|
"""A wrapper around monodepth network to extract features."""
|
|
|
|
def __init__(
|
|
self,
|
|
monodepth_encoder: SlidingPyramidNetwork,
|
|
output_dims: list[int] | None = None,
|
|
freeze_projection: bool = False,
|
|
) -> None:
|
|
"""Initialize MonodepthFeatureExtractor."""
|
|
super().__init__()
|
|
|
|
self.encoder = monodepth_encoder
|
|
|
|
# The monodepth network returns two feature maps for the first entry in
|
|
# backbone.encoder.dims_encoder.
|
|
monodepth_dims = self.encoder.dims_encoder
|
|
monodepth_dims = monodepth_dims
|
|
|
|
if output_dims is not None:
|
|
if not len(output_dims) == len(monodepth_dims):
|
|
raise ValueError(
|
|
"When set, number of output dimensions must be equal to output "
|
|
f"dimensions of monodepth model {len(monodepth_dims)}."
|
|
)
|
|
|
|
self.projection = ProjectionModule(monodepth_dims, output_dims)
|
|
self.output_dims = output_dims
|
|
else:
|
|
self.projection = nn.Identity()
|
|
self.output_dims = monodepth_dims
|
|
|
|
if freeze_projection:
|
|
self.projection.requires_grad_(False)
|
|
|
|
def forward(self, input_features: torch.Tensor) -> list[torch.Tensor]:
|
|
"""Extract multi-resolution features."""
|
|
encodings = self.encoder(input_features[:, :3].contiguous())
|
|
return self.projection(encodings)
|
|
|
|
def internal_resolution(self) -> int:
|
|
"""Internal resolution of the encoder."""
|
|
return self.encoder.internal_resolution()
|