268 lines
9.0 KiB
Python
268 lines
9.0 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
|
|
|
|
from typing import NamedTuple
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from sharp.models.blocks import (
|
|
FeatureFusionBlock2d,
|
|
NormLayerName,
|
|
residual_block_2d,
|
|
)
|
|
from sharp.models.decoders import BaseDecoder, MultiresConvDecoder
|
|
from sharp.models.params import DPTImageEncoderType, GaussianDecoderParams
|
|
|
|
|
|
def create_gaussian_decoder(
|
|
params: GaussianDecoderParams, dims_depth_features: list[int]
|
|
) -> GaussianDensePredictionTransformer:
|
|
"""Create gaussian_decoder model specified by gaussian_decoder_name."""
|
|
decoder = MultiresConvDecoder(
|
|
dims_depth_features,
|
|
params.dims_decoder,
|
|
grad_checkpointing=params.grad_checkpointing,
|
|
upsampling_mode=params.upsampling_mode,
|
|
)
|
|
|
|
return GaussianDensePredictionTransformer(
|
|
decoder=decoder,
|
|
dim_in=params.dim_in,
|
|
dim_out=params.dim_out,
|
|
stride_out=params.stride,
|
|
norm_type=params.norm_type,
|
|
norm_num_groups=params.norm_num_groups,
|
|
use_depth_input=params.use_depth_input,
|
|
grad_checkpointing=params.grad_checkpointing,
|
|
image_encoder_type=params.image_encoder_type,
|
|
image_encoder_params=params,
|
|
)
|
|
|
|
|
|
def _create_project_upsample_block(
|
|
dim_in: int,
|
|
dim_out: int,
|
|
upsample_layers: int,
|
|
dim_intermediate: int | None = None,
|
|
) -> nn.Module:
|
|
if dim_intermediate is None:
|
|
dim_intermediate = dim_out
|
|
# Projection.
|
|
blocks = [
|
|
nn.Conv2d(
|
|
in_channels=dim_in,
|
|
out_channels=dim_intermediate,
|
|
kernel_size=1,
|
|
stride=1,
|
|
padding=0,
|
|
bias=False,
|
|
)
|
|
]
|
|
|
|
# Upsampling.
|
|
blocks += [
|
|
nn.ConvTranspose2d(
|
|
in_channels=dim_intermediate if i == 0 else dim_out,
|
|
out_channels=dim_out,
|
|
kernel_size=2,
|
|
stride=2,
|
|
padding=0,
|
|
bias=False,
|
|
)
|
|
for i in range(upsample_layers)
|
|
]
|
|
|
|
return nn.Sequential(*blocks)
|
|
|
|
|
|
class ImageFeatures(NamedTuple):
|
|
"""Image feature extracted from decoder."""
|
|
|
|
texture_features: torch.Tensor
|
|
geometry_features: torch.Tensor
|
|
|
|
|
|
class SkipConvBackbone(nn.Module):
|
|
"""A wrapper around a conv layer that behaves like a BaseBackbone."""
|
|
|
|
def __init__(self, dim_in: int, dim_out: int, kernel_size: int, stride_out: int):
|
|
"""Initialize SkipConvBackbone."""
|
|
super().__init__()
|
|
self.stride_out = stride_out
|
|
if stride_out == 1 and kernel_size != 1:
|
|
raise ValueError("We only support kernel_size = 1 if stride_out is 1.")
|
|
padding: int = (kernel_size - 1) // 2
|
|
self.conv = nn.Conv2d(
|
|
dim_in, dim_out, kernel_size=kernel_size, stride=stride_out, padding=padding
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
input_features: torch.Tensor,
|
|
encodings: list[torch.Tensor] | None = None,
|
|
) -> ImageFeatures:
|
|
"""Apply SkipConvBackbone to image."""
|
|
output = self.conv(input_features)
|
|
return ImageFeatures(
|
|
texture_features=output,
|
|
geometry_features=output,
|
|
)
|
|
|
|
@property
|
|
def stride(self) -> int:
|
|
"""Effective downsampling stride."""
|
|
return self.stride_out
|
|
|
|
|
|
class GaussianDensePredictionTransformer(nn.Module):
|
|
"""Dense Prediction Transformer for Gaussian.
|
|
|
|
Reuse monodepth decoded features for processing.
|
|
"""
|
|
|
|
norm_type: NormLayerName
|
|
|
|
def __init__(
|
|
self,
|
|
decoder: BaseDecoder,
|
|
dim_in: int,
|
|
dim_out: int,
|
|
stride_out: int,
|
|
image_encoder_params: GaussianDecoderParams,
|
|
image_encoder_type: DPTImageEncoderType = "skip_conv",
|
|
norm_type: NormLayerName = "group_norm",
|
|
norm_num_groups: int = 8,
|
|
use_depth_input: bool = True,
|
|
grad_checkpointing: bool = False,
|
|
):
|
|
"""Initialize Dense Prediction Transformer for Gaussian.
|
|
|
|
Args:
|
|
decoder: Decoder to decode features.
|
|
monodepth_decoder: Optional monodepth decoder to fuse monodepth decoded features.
|
|
dim_in: Input dimension.
|
|
dim_out: Final output dimension.
|
|
stride_out: Stride of output feature map.
|
|
image_encoder_params: The backbone parameters to configurate the image encoder.
|
|
image_encoder_type: Type of image encoder to use.
|
|
encoder: Encoder to generate features using monodepth model.
|
|
norm_type: Type of norm layers.
|
|
norm_num_groups: Num groups for norm layers.
|
|
use_depth_input: Whether to use depth input.
|
|
grad_checkpointing: Whether to use gradient checkpointing.
|
|
"""
|
|
super().__init__()
|
|
|
|
self.decoder = decoder
|
|
self.dim_in = dim_in
|
|
self.dim_out = dim_out
|
|
self.stride_out = stride_out
|
|
self.norm_type = norm_type
|
|
self.norm_num_groups = norm_num_groups
|
|
self.use_depth_input = use_depth_input
|
|
self.grad_checkpointing = grad_checkpointing
|
|
self.image_encoder_type = image_encoder_type
|
|
|
|
# Adopt an image encoder to lift dimension to monodepth feature and
|
|
# resize to be the same resolution as the decoder output.
|
|
dim_in = self.dim_in if use_depth_input else self.dim_in - 1
|
|
image_encoder_params.dim_in = dim_in
|
|
image_encoder_params.dim_out = decoder.dim_out
|
|
self.image_encoder = self._create_image_encoder(image_encoder_params, stride_out)
|
|
|
|
self.fusion = FeatureFusionBlock2d(decoder.dim_out)
|
|
|
|
if stride_out == 1:
|
|
self.upsample = _create_project_upsample_block(
|
|
decoder.dim_out,
|
|
decoder.dim_out,
|
|
upsample_layers=1,
|
|
)
|
|
elif stride_out == 2:
|
|
self.upsample = nn.Identity()
|
|
else:
|
|
raise ValueError("We only support stride is 1 or 2 for DPT backbone.")
|
|
|
|
self.texture_head = self._create_head(dim_decoder=decoder.dim_out, dim_out=self.dim_out)
|
|
self.geometry_head = self._create_head(dim_decoder=decoder.dim_out, dim_out=self.dim_out)
|
|
|
|
def _create_head(self, dim_decoder: int, dim_out: int) -> nn.Module:
|
|
return nn.Sequential(
|
|
residual_block_2d(
|
|
dim_in=dim_decoder,
|
|
dim_out=dim_decoder,
|
|
dim_hidden=dim_decoder // 2,
|
|
norm_type=self.norm_type,
|
|
norm_num_groups=self.norm_num_groups,
|
|
),
|
|
residual_block_2d(
|
|
dim_in=dim_decoder,
|
|
dim_hidden=dim_decoder // 2,
|
|
dim_out=dim_decoder,
|
|
norm_type=self.norm_type,
|
|
norm_num_groups=self.norm_num_groups,
|
|
),
|
|
nn.ReLU(),
|
|
nn.Conv2d(dim_decoder, dim_out, kernel_size=1, stride=1),
|
|
nn.ReLU(),
|
|
)
|
|
|
|
def _create_image_encoder(
|
|
self, image_encoder_params: GaussianDecoderParams, stride_out: int
|
|
) -> nn.Module:
|
|
"""Create image encoder and return based on parameters."""
|
|
if self.image_encoder_type == "skip_conv":
|
|
# Use kernel_size = 1 only if stride_out is 1.
|
|
return SkipConvBackbone(
|
|
image_encoder_params.dim_in,
|
|
image_encoder_params.dim_out,
|
|
kernel_size=3 if stride_out != 1 else 1,
|
|
stride_out=stride_out,
|
|
)
|
|
elif self.image_encoder_type == "skip_conv_kernel2":
|
|
return SkipConvBackbone(
|
|
image_encoder_params.dim_in,
|
|
image_encoder_params.dim_out,
|
|
kernel_size=stride_out,
|
|
stride_out=stride_out,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported image encoder type: {self.image_encoder_type}")
|
|
|
|
def forward(self, input_features: torch.Tensor, encodings: list[torch.Tensor]) -> ImageFeatures:
|
|
"""Run monodepth and fuse features with input image to predict Gaussians.
|
|
|
|
Args:
|
|
input_features: The input features to use.
|
|
encodings: Feature encodings (e.g. from monodepth network).
|
|
"""
|
|
features = self.decoder(encodings).contiguous()
|
|
features = self.upsample(features)
|
|
|
|
if self.use_depth_input:
|
|
skip_features = self.image_encoder(input_features).texture_features
|
|
else:
|
|
skip_features = self.image_encoder(input_features[:, :3].contiguous())
|
|
features = self.fusion(features, skip_features)
|
|
|
|
texture_features = self.texture_head(features)
|
|
geometry_features = self.geometry_head(features)
|
|
|
|
return ImageFeatures(
|
|
texture_features=texture_features, # type: ignore
|
|
geometry_features=geometry_features, # type: ignore
|
|
)
|
|
|
|
@property
|
|
def stride(self) -> int:
|
|
"""Internal stride of GaussianDensePredictionTransformer."""
|
|
return self.stride_out
|