269 lines
9.6 KiB
Python
269 lines
9.6 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 copy
|
|
from typing import NamedTuple, Tuple
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from sharp.models import normalizers
|
|
from sharp.models.decoders import MultiresConvDecoder, create_monodepth_decoder
|
|
from sharp.models.encoders import (
|
|
SlidingPyramidNetwork,
|
|
create_monodepth_encoder,
|
|
)
|
|
from sharp.utils import module_surgery
|
|
|
|
from .params import MonodepthAdaptorParams, MonodepthParams
|
|
|
|
DimsDecoder = Tuple[int, int, int, int, int]
|
|
|
|
|
|
class MonodepthDensePredictionTransformer(nn.Module):
|
|
"""Dense Prediction Transformer for monodepth.
|
|
|
|
Attach the disparity prediction head for monodepth prediction.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
encoder: SlidingPyramidNetwork,
|
|
decoder: MultiresConvDecoder,
|
|
last_dims: tuple[int, int],
|
|
):
|
|
"""Initialize Dense Prediction Transformer.
|
|
|
|
Args:
|
|
encoder: The SlidingPyramidTransformer backbone.
|
|
decoder: The MultiresConvDecoder decoder.
|
|
last_dims: The dimension for the last convolution layers.
|
|
"""
|
|
super().__init__()
|
|
|
|
self.normalizer = normalizers.AffineRangeNormalizer(
|
|
input_range=(0, 1), output_range=(-1, 1)
|
|
)
|
|
self.encoder = encoder
|
|
self.decoder = decoder
|
|
|
|
dim_decoder = decoder.dim_out
|
|
self.head = nn.Sequential(
|
|
nn.Conv2d(dim_decoder, dim_decoder // 2, kernel_size=3, stride=1, padding=1),
|
|
nn.ConvTranspose2d(
|
|
in_channels=dim_decoder // 2,
|
|
out_channels=dim_decoder // 2,
|
|
kernel_size=2,
|
|
stride=2,
|
|
padding=0,
|
|
bias=True,
|
|
),
|
|
nn.Conv2d(
|
|
dim_decoder // 2,
|
|
last_dims[0],
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
),
|
|
nn.ReLU(True),
|
|
nn.Conv2d(last_dims[0], last_dims[1], kernel_size=1, stride=1, padding=0),
|
|
nn.ReLU(),
|
|
)
|
|
|
|
# Set the final convoultion layer's bias to be 0.
|
|
self.head[4].bias.data.fill_(0)
|
|
|
|
self.grad_checkpointing = False
|
|
|
|
@torch.jit.ignore
|
|
def set_grad_checkpointing(self, is_enabled=True):
|
|
"""Enable grad checkpointing."""
|
|
self.grad_checkpointing = is_enabled
|
|
self.encoder.set_grad_checkpointing(self.grad_checkpointing)
|
|
self.decoder.set_grad_checkpointing(self.grad_checkpointing)
|
|
|
|
def forward(self, image: torch.Tensor) -> torch.Tensor:
|
|
"""Decode by projection and fusion of multi-resolution encodings."""
|
|
encodings = self.encoder(self.normalizer(image))
|
|
num_encoder_features = len(self.encoder.dims_encoder)
|
|
features = self.decoder(encodings[:num_encoder_features])
|
|
disparity = self.head(features)
|
|
return disparity
|
|
|
|
def internal_resolution(self) -> int:
|
|
"""Return the internal image size of the network."""
|
|
return self.encoder.internal_resolution()
|
|
|
|
|
|
def create_monodepth_dpt(
|
|
params: MonodepthParams | None = None,
|
|
) -> MonodepthDensePredictionTransformer:
|
|
"""Creates DepthDensePredictionTransformer model.
|
|
|
|
Args:
|
|
params: Parameters of monodepth network.
|
|
|
|
Returns:
|
|
The configured monodepth DPT.
|
|
"""
|
|
if params is None:
|
|
params = MonodepthParams()
|
|
encoder: SlidingPyramidNetwork = create_monodepth_encoder(
|
|
params.patch_encoder_preset,
|
|
params.image_encoder_preset,
|
|
use_patch_overlap=params.use_patch_overlap,
|
|
last_encoder=params.dims_decoder[0],
|
|
)
|
|
|
|
decoder: MultiresConvDecoder = create_monodepth_decoder(
|
|
params.patch_encoder_preset, params.dims_decoder
|
|
)
|
|
|
|
monodepth_model = MonodepthDensePredictionTransformer(
|
|
encoder=encoder, decoder=decoder, last_dims=(32, 1)
|
|
)
|
|
|
|
# By default, we don't train the monodepth model.
|
|
# However, we allow to selectively unfreeze parts of the network.
|
|
monodepth_model.requires_grad_(False)
|
|
|
|
monodepth_model.encoder.set_requires_grad_(
|
|
patch_encoder=params.unfreeze_patch_encoder,
|
|
image_encoder=params.unfreeze_image_encoder,
|
|
)
|
|
monodepth_model.decoder.requires_grad_(params.unfreeze_decoder)
|
|
monodepth_model.head.requires_grad_(params.unfreeze_head)
|
|
|
|
if not params.unfreeze_norm_layers:
|
|
module_surgery.freeze_norm_layer(monodepth_model)
|
|
|
|
monodepth_model.set_grad_checkpointing(params.grad_checkpointing)
|
|
|
|
return monodepth_model
|
|
|
|
|
|
class MonodepthOutput(NamedTuple):
|
|
"""Output of the monodepth model."""
|
|
|
|
# Disparity output from the monodepth model.
|
|
disparity: torch.Tensor
|
|
# Multi-level features from monodepth encoder.
|
|
encoder_features: list[torch.Tensor]
|
|
# Single-level feature from monodepth decoder.
|
|
decoder_features: torch.Tensor
|
|
# List of monodepth features to be used in gaussian predictor.
|
|
output_features: list[torch.Tensor]
|
|
# List of intermediate encoder features to be used in distillation.
|
|
intermediate_features: list[torch.Tensor] = []
|
|
|
|
|
|
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,
|
|
output_features=output_features,
|
|
intermediate_features=intermediate_features,
|
|
)
|
|
|
|
def get_feature_dims(self) -> list[int]:
|
|
"""Return dimensions of output feature maps."""
|
|
dims = []
|
|
if self.return_encoder_features:
|
|
dims.extend(self.monodepth_predictor.encoder.dims_encoder)
|
|
|
|
if self.return_decoder_features:
|
|
dims.append(self.monodepth_predictor.decoder.dim_out)
|
|
|
|
return dims
|
|
|
|
def internal_resolution(self) -> int:
|
|
"""Return the internal image size of the network."""
|
|
return self.monodepth_predictor.internal_resolution()
|
|
|
|
def replicate_head(self, num_repeat: int):
|
|
"""Replicate the last convolution layer (head[4] in DPT) for multi layer depth."""
|
|
conv_last = copy.deepcopy(self.monodepth_predictor.head[4])
|
|
self.monodepth_predictor.head[4].out_channels = num_repeat
|
|
self.monodepth_predictor.head[4].weight = nn.Parameter(
|
|
conv_last.weight.repeat(num_repeat, 1, 1, 1)
|
|
)
|
|
self.monodepth_predictor.head[4].bias = nn.Parameter(conv_last.bias.repeat(num_repeat))
|
|
|
|
|
|
def create_monodepth_adaptor(
|
|
monodepth_predictor: MonodepthDensePredictionTransformer,
|
|
params: MonodepthAdaptorParams,
|
|
num_monodepth_layers: int,
|
|
sorting_monodepth: bool,
|
|
) -> MonodepthWithEncodingAdaptor:
|
|
"""Create an adaptor that returns both disparity and features."""
|
|
adaptor = MonodepthWithEncodingAdaptor(
|
|
monodepth_predictor=monodepth_predictor,
|
|
return_encoder_features=params.encoder_features,
|
|
return_decoder_features=params.decoder_features,
|
|
num_monodepth_layers=num_monodepth_layers,
|
|
sorting_monodepth=sorting_monodepth,
|
|
)
|
|
return adaptor
|