vendored model
This commit is contained in:
Andrea Pozzetti
@@ -0,0 +1,267 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user