Spaces:
Sleeping
Sleeping
# Copyright (c) Meta Platforms, Inc. and affiliates. | |
# | |
# This source code is licensed under the Apache License, Version 2.0 | |
# found in the LICENSE file in the root directory of this source tree. | |
import copy | |
from functools import partial | |
import math | |
import warnings | |
import torch | |
import torch.nn as nn | |
from .ops import resize | |
# XXX: (Untested) replacement for mmcv.imdenormalize() | |
def _imdenormalize(img, mean, std, to_bgr=True): | |
import numpy as np | |
mean = mean.reshape(1, -1).astype(np.float64) | |
std = std.reshape(1, -1).astype(np.float64) | |
img = (img * std) + mean | |
if to_bgr: | |
img = img[::-1] | |
return img | |
class DepthBaseDecodeHead(nn.Module): | |
"""Base class for BaseDecodeHead. | |
Args: | |
in_channels (List): Input channels. | |
channels (int): Channels after modules, before conv_depth. | |
conv_layer (nn.Module): Conv layers. Default: None. | |
act_layer (nn.Module): Activation layers. Default: nn.ReLU. | |
loss_decode (dict): Config of decode loss. | |
Default: (). | |
sampler (dict|None): The config of depth map sampler. | |
Default: None. | |
align_corners (bool): align_corners argument of F.interpolate. | |
Default: False. | |
min_depth (int): Min depth in dataset setting. | |
Default: 1e-3. | |
max_depth (int): Max depth in dataset setting. | |
Default: None. | |
norm_layer (dict|None): Norm layers. | |
Default: None. | |
classify (bool): Whether predict depth in a cls.-reg. manner. | |
Default: False. | |
n_bins (int): The number of bins used in cls. step. | |
Default: 256. | |
bins_strategy (str): The discrete strategy used in cls. step. | |
Default: 'UD'. | |
norm_strategy (str): The norm strategy on cls. probability | |
distribution. Default: 'linear' | |
scale_up (str): Whether predict depth in a scale-up manner. | |
Default: False. | |
""" | |
def __init__( | |
self, | |
in_channels, | |
conv_layer=None, | |
act_layer=nn.ReLU, | |
channels=96, | |
loss_decode=(), | |
sampler=None, | |
align_corners=False, | |
min_depth=1e-3, | |
max_depth=None, | |
norm_layer=None, | |
classify=False, | |
n_bins=256, | |
bins_strategy="UD", | |
norm_strategy="linear", | |
scale_up=False, | |
): | |
super(DepthBaseDecodeHead, self).__init__() | |
self.in_channels = in_channels | |
self.channels = channels | |
self.conf_layer = conv_layer | |
self.act_layer = act_layer | |
self.loss_decode = loss_decode | |
self.align_corners = align_corners | |
self.min_depth = min_depth | |
self.max_depth = max_depth | |
self.norm_layer = norm_layer | |
self.classify = classify | |
self.n_bins = n_bins | |
self.scale_up = scale_up | |
if self.classify: | |
assert bins_strategy in ["UD", "SID"], "Support bins_strategy: UD, SID" | |
assert norm_strategy in ["linear", "softmax", "sigmoid"], "Support norm_strategy: linear, softmax, sigmoid" | |
self.bins_strategy = bins_strategy | |
self.norm_strategy = norm_strategy | |
self.softmax = nn.Softmax(dim=1) | |
self.conv_depth = nn.Conv2d(channels, n_bins, kernel_size=3, padding=1, stride=1) | |
else: | |
self.conv_depth = nn.Conv2d(channels, 1, kernel_size=3, padding=1, stride=1) | |
self.relu = nn.ReLU() | |
self.sigmoid = nn.Sigmoid() | |
def forward(self, inputs, img_metas): | |
"""Placeholder of forward function.""" | |
pass | |
def forward_train(self, img, inputs, img_metas, depth_gt): | |
"""Forward function for training. | |
Args: | |
inputs (list[Tensor]): List of multi-level img features. | |
img_metas (list[dict]): List of image info dict where each dict | |
has: 'img_shape', 'scale_factor', 'flip', and may also contain | |
'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'. | |
For details on the values of these keys see | |
`depth/datasets/pipelines/formatting.py:Collect`. | |
depth_gt (Tensor): GT depth | |
Returns: | |
dict[str, Tensor]: a dictionary of loss components | |
""" | |
depth_pred = self.forward(inputs, img_metas) | |
losses = self.losses(depth_pred, depth_gt) | |
log_imgs = self.log_images(img[0], depth_pred[0], depth_gt[0], img_metas[0]) | |
losses.update(**log_imgs) | |
return losses | |
def forward_test(self, inputs, img_metas): | |
"""Forward function for testing. | |
Args: | |
inputs (list[Tensor]): List of multi-level img features. | |
img_metas (list[dict]): List of image info dict where each dict | |
has: 'img_shape', 'scale_factor', 'flip', and may also contain | |
'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'. | |
For details on the values of these keys see | |
`depth/datasets/pipelines/formatting.py:Collect`. | |
Returns: | |
Tensor: Output depth map. | |
""" | |
return self.forward(inputs, img_metas) | |
def depth_pred(self, feat): | |
"""Prediction each pixel.""" | |
if self.classify: | |
logit = self.conv_depth(feat) | |
if self.bins_strategy == "UD": | |
bins = torch.linspace(self.min_depth, self.max_depth, self.n_bins, device=feat.device) | |
elif self.bins_strategy == "SID": | |
bins = torch.logspace(self.min_depth, self.max_depth, self.n_bins, device=feat.device) | |
# following Adabins, default linear | |
if self.norm_strategy == "linear": | |
logit = torch.relu(logit) | |
eps = 0.1 | |
logit = logit + eps | |
logit = logit / logit.sum(dim=1, keepdim=True) | |
elif self.norm_strategy == "softmax": | |
logit = torch.softmax(logit, dim=1) | |
elif self.norm_strategy == "sigmoid": | |
logit = torch.sigmoid(logit) | |
logit = logit / logit.sum(dim=1, keepdim=True) | |
output = torch.einsum("ikmn,k->imn", [logit, bins]).unsqueeze(dim=1) | |
else: | |
if self.scale_up: | |
output = self.sigmoid(self.conv_depth(feat)) * self.max_depth | |
else: | |
output = self.relu(self.conv_depth(feat)) + self.min_depth | |
return output | |
def losses(self, depth_pred, depth_gt): | |
"""Compute depth loss.""" | |
loss = dict() | |
depth_pred = resize( | |
input=depth_pred, size=depth_gt.shape[2:], mode="bilinear", align_corners=self.align_corners, warning=False | |
) | |
if not isinstance(self.loss_decode, nn.ModuleList): | |
losses_decode = [self.loss_decode] | |
else: | |
losses_decode = self.loss_decode | |
for loss_decode in losses_decode: | |
if loss_decode.loss_name not in loss: | |
loss[loss_decode.loss_name] = loss_decode(depth_pred, depth_gt) | |
else: | |
loss[loss_decode.loss_name] += loss_decode(depth_pred, depth_gt) | |
return loss | |
def log_images(self, img_path, depth_pred, depth_gt, img_meta): | |
import numpy as np | |
show_img = copy.deepcopy(img_path.detach().cpu().permute(1, 2, 0)) | |
show_img = show_img.numpy().astype(np.float32) | |
show_img = _imdenormalize( | |
show_img, | |
img_meta["img_norm_cfg"]["mean"], | |
img_meta["img_norm_cfg"]["std"], | |
img_meta["img_norm_cfg"]["to_rgb"], | |
) | |
show_img = np.clip(show_img, 0, 255) | |
show_img = show_img.astype(np.uint8) | |
show_img = show_img[:, :, ::-1] | |
show_img = show_img.transpose(0, 2, 1) | |
show_img = show_img.transpose(1, 0, 2) | |
depth_pred = depth_pred / torch.max(depth_pred) | |
depth_gt = depth_gt / torch.max(depth_gt) | |
depth_pred_color = copy.deepcopy(depth_pred.detach().cpu()) | |
depth_gt_color = copy.deepcopy(depth_gt.detach().cpu()) | |
return {"img_rgb": show_img, "img_depth_pred": depth_pred_color, "img_depth_gt": depth_gt_color} | |
class BNHead(DepthBaseDecodeHead): | |
"""Just a batchnorm.""" | |
def __init__(self, input_transform="resize_concat", in_index=(0, 1, 2, 3), upsample=1, **kwargs): | |
super().__init__(**kwargs) | |
self.input_transform = input_transform | |
self.in_index = in_index | |
self.upsample = upsample | |
# self.bn = nn.SyncBatchNorm(self.in_channels) | |
if self.classify: | |
self.conv_depth = nn.Conv2d(self.channels, self.n_bins, kernel_size=1, padding=0, stride=1) | |
else: | |
self.conv_depth = nn.Conv2d(self.channels, 1, kernel_size=1, padding=0, stride=1) | |
def _transform_inputs(self, inputs): | |
"""Transform inputs for decoder. | |
Args: | |
inputs (list[Tensor]): List of multi-level img features. | |
Returns: | |
Tensor: The transformed inputs | |
""" | |
if "concat" in self.input_transform: | |
inputs = [inputs[i] for i in self.in_index] | |
if "resize" in self.input_transform: | |
inputs = [ | |
resize( | |
input=x, | |
size=[s * self.upsample for s in inputs[0].shape[2:]], | |
mode="bilinear", | |
align_corners=self.align_corners, | |
) | |
for x in inputs | |
] | |
inputs = torch.cat(inputs, dim=1) | |
elif self.input_transform == "multiple_select": | |
inputs = [inputs[i] for i in self.in_index] | |
else: | |
inputs = inputs[self.in_index] | |
return inputs | |
def _forward_feature(self, inputs, img_metas=None, **kwargs): | |
"""Forward function for feature maps before classifying each pixel with | |
``self.cls_seg`` fc. | |
Args: | |
inputs (list[Tensor]): List of multi-level img features. | |
Returns: | |
feats (Tensor): A tensor of shape (batch_size, self.channels, | |
H, W) which is feature map for last layer of decoder head. | |
""" | |
# accept lists (for cls token) | |
inputs = list(inputs) | |
for i, x in enumerate(inputs): | |
if len(x) == 2: | |
x, cls_token = x[0], x[1] | |
if len(x.shape) == 2: | |
x = x[:, :, None, None] | |
cls_token = cls_token[:, :, None, None].expand_as(x) | |
inputs[i] = torch.cat((x, cls_token), 1) | |
else: | |
x = x[0] | |
if len(x.shape) == 2: | |
x = x[:, :, None, None] | |
inputs[i] = x | |
x = self._transform_inputs(inputs) | |
# feats = self.bn(x) | |
return x | |
def forward(self, inputs, img_metas=None, **kwargs): | |
"""Forward function.""" | |
output = self._forward_feature(inputs, img_metas=img_metas, **kwargs) | |
output = self.depth_pred(output) | |
return output | |
class ConvModule(nn.Module): | |
"""A conv block that bundles conv/norm/activation layers. | |
This block simplifies the usage of convolution layers, which are commonly | |
used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU). | |
It is based upon three build methods: `build_conv_layer()`, | |
`build_norm_layer()` and `build_activation_layer()`. | |
Besides, we add some additional features in this module. | |
1. Automatically set `bias` of the conv layer. | |
2. Spectral norm is supported. | |
3. More padding modes are supported. Before PyTorch 1.5, nn.Conv2d only | |
supports zero and circular padding, and we add "reflect" padding mode. | |
Args: | |
in_channels (int): Number of channels in the input feature map. | |
Same as that in ``nn._ConvNd``. | |
out_channels (int): Number of channels produced by the convolution. | |
Same as that in ``nn._ConvNd``. | |
kernel_size (int | tuple[int]): Size of the convolving kernel. | |
Same as that in ``nn._ConvNd``. | |
stride (int | tuple[int]): Stride of the convolution. | |
Same as that in ``nn._ConvNd``. | |
padding (int | tuple[int]): Zero-padding added to both sides of | |
the input. Same as that in ``nn._ConvNd``. | |
dilation (int | tuple[int]): Spacing between kernel elements. | |
Same as that in ``nn._ConvNd``. | |
groups (int): Number of blocked connections from input channels to | |
output channels. Same as that in ``nn._ConvNd``. | |
bias (bool | str): If specified as `auto`, it will be decided by the | |
norm_layer. Bias will be set as True if `norm_layer` is None, otherwise | |
False. Default: "auto". | |
conv_layer (nn.Module): Convolution layer. Default: None, | |
which means using conv2d. | |
norm_layer (nn.Module): Normalization layer. Default: None. | |
act_layer (nn.Module): Activation layer. Default: nn.ReLU. | |
inplace (bool): Whether to use inplace mode for activation. | |
Default: True. | |
with_spectral_norm (bool): Whether use spectral norm in conv module. | |
Default: False. | |
padding_mode (str): If the `padding_mode` has not been supported by | |
current `Conv2d` in PyTorch, we will use our own padding layer | |
instead. Currently, we support ['zeros', 'circular'] with official | |
implementation and ['reflect'] with our own implementation. | |
Default: 'zeros'. | |
order (tuple[str]): The order of conv/norm/activation layers. It is a | |
sequence of "conv", "norm" and "act". Common examples are | |
("conv", "norm", "act") and ("act", "conv", "norm"). | |
Default: ('conv', 'norm', 'act'). | |
""" | |
_abbr_ = "conv_block" | |
def __init__( | |
self, | |
in_channels, | |
out_channels, | |
kernel_size, | |
stride=1, | |
padding=0, | |
dilation=1, | |
groups=1, | |
bias="auto", | |
conv_layer=nn.Conv2d, | |
norm_layer=None, | |
act_layer=nn.ReLU, | |
inplace=True, | |
with_spectral_norm=False, | |
padding_mode="zeros", | |
order=("conv", "norm", "act"), | |
): | |
super(ConvModule, self).__init__() | |
official_padding_mode = ["zeros", "circular"] | |
self.conv_layer = conv_layer | |
self.norm_layer = norm_layer | |
self.act_layer = act_layer | |
self.inplace = inplace | |
self.with_spectral_norm = with_spectral_norm | |
self.with_explicit_padding = padding_mode not in official_padding_mode | |
self.order = order | |
assert isinstance(self.order, tuple) and len(self.order) == 3 | |
assert set(order) == set(["conv", "norm", "act"]) | |
self.with_norm = norm_layer is not None | |
self.with_activation = act_layer is not None | |
# if the conv layer is before a norm layer, bias is unnecessary. | |
if bias == "auto": | |
bias = not self.with_norm | |
self.with_bias = bias | |
if self.with_explicit_padding: | |
if padding_mode == "zeros": | |
padding_layer = nn.ZeroPad2d | |
else: | |
raise AssertionError(f"Unsupported padding mode: {padding_mode}") | |
self.pad = padding_layer(padding) | |
# reset padding to 0 for conv module | |
conv_padding = 0 if self.with_explicit_padding else padding | |
# build convolution layer | |
self.conv = self.conv_layer( | |
in_channels, | |
out_channels, | |
kernel_size, | |
stride=stride, | |
padding=conv_padding, | |
dilation=dilation, | |
groups=groups, | |
bias=bias, | |
) | |
# export the attributes of self.conv to a higher level for convenience | |
self.in_channels = self.conv.in_channels | |
self.out_channels = self.conv.out_channels | |
self.kernel_size = self.conv.kernel_size | |
self.stride = self.conv.stride | |
self.padding = padding | |
self.dilation = self.conv.dilation | |
self.transposed = self.conv.transposed | |
self.output_padding = self.conv.output_padding | |
self.groups = self.conv.groups | |
if self.with_spectral_norm: | |
self.conv = nn.utils.spectral_norm(self.conv) | |
# build normalization layers | |
if self.with_norm: | |
# norm layer is after conv layer | |
if order.index("norm") > order.index("conv"): | |
norm_channels = out_channels | |
else: | |
norm_channels = in_channels | |
norm = partial(norm_layer, num_features=norm_channels) | |
self.add_module("norm", norm) | |
if self.with_bias: | |
from torch.nnModules.batchnorm import _BatchNorm | |
from torch.nnModules.instancenorm import _InstanceNorm | |
if isinstance(norm, (_BatchNorm, _InstanceNorm)): | |
warnings.warn("Unnecessary conv bias before batch/instance norm") | |
else: | |
self.norm_name = None | |
# build activation layer | |
if self.with_activation: | |
# nn.Tanh has no 'inplace' argument | |
# (nn.Tanh, nn.PReLU, nn.Sigmoid, nn.HSigmoid, nn.Swish, nn.GELU) | |
if not isinstance(act_layer, (nn.Tanh, nn.PReLU, nn.Sigmoid, nn.GELU)): | |
act_layer = partial(act_layer, inplace=inplace) | |
self.activate = act_layer() | |
# Use msra init by default | |
self.init_weights() | |
def norm(self): | |
if self.norm_name: | |
return getattr(self, self.norm_name) | |
else: | |
return None | |
def init_weights(self): | |
# 1. It is mainly for customized conv layers with their own | |
# initialization manners by calling their own ``init_weights()``, | |
# and we do not want ConvModule to override the initialization. | |
# 2. For customized conv layers without their own initialization | |
# manners (that is, they don't have their own ``init_weights()``) | |
# and PyTorch's conv layers, they will be initialized by | |
# this method with default ``kaiming_init``. | |
# Note: For PyTorch's conv layers, they will be overwritten by our | |
# initialization implementation using default ``kaiming_init``. | |
if not hasattr(self.conv, "init_weights"): | |
if self.with_activation and isinstance(self.act_layer, nn.LeakyReLU): | |
nonlinearity = "leaky_relu" | |
a = 0.01 # XXX: default negative_slope | |
else: | |
nonlinearity = "relu" | |
a = 0 | |
if hasattr(self.conv, "weight") and self.conv.weight is not None: | |
nn.init.kaiming_normal_(self.conv.weight, a=a, mode="fan_out", nonlinearity=nonlinearity) | |
if hasattr(self.conv, "bias") and self.conv.bias is not None: | |
nn.init.constant_(self.conv.bias, 0) | |
if self.with_norm: | |
if hasattr(self.norm, "weight") and self.norm.weight is not None: | |
nn.init.constant_(self.norm.weight, 1) | |
if hasattr(self.norm, "bias") and self.norm.bias is not None: | |
nn.init.constant_(self.norm.bias, 0) | |
def forward(self, x, activate=True, norm=True): | |
for layer in self.order: | |
if layer == "conv": | |
if self.with_explicit_padding: | |
x = self.pad(x) | |
x = self.conv(x) | |
elif layer == "norm" and norm and self.with_norm: | |
x = self.norm(x) | |
elif layer == "act" and activate and self.with_activation: | |
x = self.activate(x) | |
return x | |
class Interpolate(nn.Module): | |
def __init__(self, scale_factor, mode, align_corners=False): | |
super(Interpolate, self).__init__() | |
self.interp = nn.functional.interpolate | |
self.scale_factor = scale_factor | |
self.mode = mode | |
self.align_corners = align_corners | |
def forward(self, x): | |
x = self.interp(x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners) | |
return x | |
class HeadDepth(nn.Module): | |
def __init__(self, features): | |
super(HeadDepth, self).__init__() | |
self.head = nn.Sequential( | |
nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1), | |
Interpolate(scale_factor=2, mode="bilinear", align_corners=True), | |
nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1), | |
nn.ReLU(), | |
nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0), | |
) | |
def forward(self, x): | |
x = self.head(x) | |
return x | |
class ReassembleBlocks(nn.Module): | |
"""ViTPostProcessBlock, process cls_token in ViT backbone output and | |
rearrange the feature vector to feature map. | |
Args: | |
in_channels (int): ViT feature channels. Default: 768. | |
out_channels (List): output channels of each stage. | |
Default: [96, 192, 384, 768]. | |
readout_type (str): Type of readout operation. Default: 'ignore'. | |
patch_size (int): The patch size. Default: 16. | |
""" | |
def __init__(self, in_channels=768, out_channels=[96, 192, 384, 768], readout_type="ignore", patch_size=16): | |
super(ReassembleBlocks, self).__init__() | |
assert readout_type in ["ignore", "add", "project"] | |
self.readout_type = readout_type | |
self.patch_size = patch_size | |
self.projects = nn.ModuleList( | |
[ | |
ConvModule( | |
in_channels=in_channels, | |
out_channels=out_channel, | |
kernel_size=1, | |
act_layer=None, | |
) | |
for out_channel in out_channels | |
] | |
) | |
self.resize_layers = nn.ModuleList( | |
[ | |
nn.ConvTranspose2d( | |
in_channels=out_channels[0], out_channels=out_channels[0], kernel_size=4, stride=4, padding=0 | |
), | |
nn.ConvTranspose2d( | |
in_channels=out_channels[1], out_channels=out_channels[1], kernel_size=2, stride=2, padding=0 | |
), | |
nn.Identity(), | |
nn.Conv2d( | |
in_channels=out_channels[3], out_channels=out_channels[3], kernel_size=3, stride=2, padding=1 | |
), | |
] | |
) | |
if self.readout_type == "project": | |
self.readout_projects = nn.ModuleList() | |
for _ in range(len(self.projects)): | |
self.readout_projects.append(nn.Sequential(nn.Linear(2 * in_channels, in_channels), nn.GELU())) | |
def forward(self, inputs): | |
assert isinstance(inputs, list) | |
out = [] | |
for i, x in enumerate(inputs): | |
assert len(x) == 2 | |
x, cls_token = x[0], x[1] | |
feature_shape = x.shape | |
if self.readout_type == "project": | |
x = x.flatten(2).permute((0, 2, 1)) | |
readout = cls_token.unsqueeze(1).expand_as(x) | |
x = self.readout_projects[i](torch.cat((x, readout), -1)) | |
x = x.permute(0, 2, 1).reshape(feature_shape) | |
elif self.readout_type == "add": | |
x = x.flatten(2) + cls_token.unsqueeze(-1) | |
x = x.reshape(feature_shape) | |
else: | |
pass | |
x = self.projects[i](x) | |
x = self.resize_layers[i](x) | |
out.append(x) | |
return out | |
class PreActResidualConvUnit(nn.Module): | |
"""ResidualConvUnit, pre-activate residual unit. | |
Args: | |
in_channels (int): number of channels in the input feature map. | |
act_layer (nn.Module): activation layer. | |
norm_layer (nn.Module): norm layer. | |
stride (int): stride of the first block. Default: 1 | |
dilation (int): dilation rate for convs layers. Default: 1. | |
""" | |
def __init__(self, in_channels, act_layer, norm_layer, stride=1, dilation=1): | |
super(PreActResidualConvUnit, self).__init__() | |
self.conv1 = ConvModule( | |
in_channels, | |
in_channels, | |
3, | |
stride=stride, | |
padding=dilation, | |
dilation=dilation, | |
norm_layer=norm_layer, | |
act_layer=act_layer, | |
bias=False, | |
order=("act", "conv", "norm"), | |
) | |
self.conv2 = ConvModule( | |
in_channels, | |
in_channels, | |
3, | |
padding=1, | |
norm_layer=norm_layer, | |
act_layer=act_layer, | |
bias=False, | |
order=("act", "conv", "norm"), | |
) | |
def forward(self, inputs): | |
inputs_ = inputs.clone() | |
x = self.conv1(inputs) | |
x = self.conv2(x) | |
return x + inputs_ | |
class FeatureFusionBlock(nn.Module): | |
"""FeatureFusionBlock, merge feature map from different stages. | |
Args: | |
in_channels (int): Input channels. | |
act_layer (nn.Module): activation layer for ResidualConvUnit. | |
norm_layer (nn.Module): normalization layer. | |
expand (bool): Whether expand the channels in post process block. | |
Default: False. | |
align_corners (bool): align_corner setting for bilinear upsample. | |
Default: True. | |
""" | |
def __init__(self, in_channels, act_layer, norm_layer, expand=False, align_corners=True): | |
super(FeatureFusionBlock, self).__init__() | |
self.in_channels = in_channels | |
self.expand = expand | |
self.align_corners = align_corners | |
self.out_channels = in_channels | |
if self.expand: | |
self.out_channels = in_channels // 2 | |
self.project = ConvModule(self.in_channels, self.out_channels, kernel_size=1, act_layer=None, bias=True) | |
self.res_conv_unit1 = PreActResidualConvUnit( | |
in_channels=self.in_channels, act_layer=act_layer, norm_layer=norm_layer | |
) | |
self.res_conv_unit2 = PreActResidualConvUnit( | |
in_channels=self.in_channels, act_layer=act_layer, norm_layer=norm_layer | |
) | |
def forward(self, *inputs): | |
x = inputs[0] | |
if len(inputs) == 2: | |
if x.shape != inputs[1].shape: | |
res = resize(inputs[1], size=(x.shape[2], x.shape[3]), mode="bilinear", align_corners=False) | |
else: | |
res = inputs[1] | |
x = x + self.res_conv_unit1(res) | |
x = self.res_conv_unit2(x) | |
x = resize(x, scale_factor=2, mode="bilinear", align_corners=self.align_corners) | |
x = self.project(x) | |
return x | |
class DPTHead(DepthBaseDecodeHead): | |
"""Vision Transformers for Dense Prediction. | |
This head is implemented of `DPT <https://arxiv.org/abs/2103.13413>`_. | |
Args: | |
embed_dims (int): The embed dimension of the ViT backbone. | |
Default: 768. | |
post_process_channels (List): Out channels of post process conv | |
layers. Default: [96, 192, 384, 768]. | |
readout_type (str): Type of readout operation. Default: 'ignore'. | |
patch_size (int): The patch size. Default: 16. | |
expand_channels (bool): Whether expand the channels in post process | |
block. Default: False. | |
""" | |
def __init__( | |
self, | |
embed_dims=768, | |
post_process_channels=[96, 192, 384, 768], | |
readout_type="ignore", | |
patch_size=16, | |
expand_channels=False, | |
**kwargs, | |
): | |
super(DPTHead, self).__init__(**kwargs) | |
self.in_channels = self.in_channels | |
self.expand_channels = expand_channels | |
self.reassemble_blocks = ReassembleBlocks(embed_dims, post_process_channels, readout_type, patch_size) | |
self.post_process_channels = [ | |
channel * math.pow(2, i) if expand_channels else channel for i, channel in enumerate(post_process_channels) | |
] | |
self.convs = nn.ModuleList() | |
for channel in self.post_process_channels: | |
self.convs.append(ConvModule(channel, self.channels, kernel_size=3, padding=1, act_layer=None, bias=False)) | |
self.fusion_blocks = nn.ModuleList() | |
for _ in range(len(self.convs)): | |
self.fusion_blocks.append(FeatureFusionBlock(self.channels, self.act_layer, self.norm_layer)) | |
self.fusion_blocks[0].res_conv_unit1 = None | |
self.project = ConvModule(self.channels, self.channels, kernel_size=3, padding=1, norm_layer=self.norm_layer) | |
self.num_fusion_blocks = len(self.fusion_blocks) | |
self.num_reassemble_blocks = len(self.reassemble_blocks.resize_layers) | |
self.num_post_process_channels = len(self.post_process_channels) | |
assert self.num_fusion_blocks == self.num_reassemble_blocks | |
assert self.num_reassemble_blocks == self.num_post_process_channels | |
self.conv_depth = HeadDepth(self.channels) | |
def forward(self, inputs, img_metas): | |
assert len(inputs) == self.num_reassemble_blocks | |
x = [inp for inp in inputs] | |
x = self.reassemble_blocks(x) | |
x = [self.convs[i](feature) for i, feature in enumerate(x)] | |
out = self.fusion_blocks[0](x[-1]) | |
for i in range(1, len(self.fusion_blocks)): | |
out = self.fusion_blocks[i](out, x[-(i + 1)]) | |
out = self.project(out) | |
out = self.depth_pred(out) | |
return out | |