first commit
This commit is contained in:
Guocheng Qian
166
midas/dpt_depth.py
Normal file
166
midas/dpt_depth.py
Normal file
@@ -0,0 +1,166 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .base_model import BaseModel
|
||||
from .blocks import (
|
||||
FeatureFusionBlock_custom,
|
||||
Interpolate,
|
||||
_make_encoder,
|
||||
forward_beit,
|
||||
forward_swin,
|
||||
forward_levit,
|
||||
forward_vit,
|
||||
)
|
||||
from .backbones.levit import stem_b4_transpose
|
||||
from timm.models.layers import get_act_layer
|
||||
|
||||
|
||||
def _make_fusion_block(features, use_bn, size = None):
|
||||
return FeatureFusionBlock_custom(
|
||||
features,
|
||||
nn.ReLU(False),
|
||||
deconv=False,
|
||||
bn=use_bn,
|
||||
expand=False,
|
||||
align_corners=True,
|
||||
size=size,
|
||||
)
|
||||
|
||||
|
||||
class DPT(BaseModel):
|
||||
def __init__(
|
||||
self,
|
||||
head,
|
||||
features=256,
|
||||
backbone="vitb_rn50_384",
|
||||
readout="project",
|
||||
channels_last=False,
|
||||
use_bn=False,
|
||||
**kwargs
|
||||
):
|
||||
|
||||
super(DPT, self).__init__()
|
||||
|
||||
self.channels_last = channels_last
|
||||
|
||||
# For the Swin, Swin 2, LeViT and Next-ViT Transformers, the hierarchical architectures prevent setting the
|
||||
# hooks freely. Instead, the hooks have to be chosen according to the ranges specified in the comments.
|
||||
hooks = {
|
||||
"beitl16_512": [5, 11, 17, 23],
|
||||
"beitl16_384": [5, 11, 17, 23],
|
||||
"beitb16_384": [2, 5, 8, 11],
|
||||
"swin2l24_384": [1, 1, 17, 1], # Allowed ranges: [0, 1], [0, 1], [ 0, 17], [ 0, 1]
|
||||
"swin2b24_384": [1, 1, 17, 1], # [0, 1], [0, 1], [ 0, 17], [ 0, 1]
|
||||
"swin2t16_256": [1, 1, 5, 1], # [0, 1], [0, 1], [ 0, 5], [ 0, 1]
|
||||
"swinl12_384": [1, 1, 17, 1], # [0, 1], [0, 1], [ 0, 17], [ 0, 1]
|
||||
"next_vit_large_6m": [2, 6, 36, 39], # [0, 2], [3, 6], [ 7, 36], [37, 39]
|
||||
"levit_384": [3, 11, 21], # [0, 3], [6, 11], [14, 21]
|
||||
"vitb_rn50_384": [0, 1, 8, 11],
|
||||
"vitb16_384": [2, 5, 8, 11],
|
||||
"vitl16_384": [5, 11, 17, 23],
|
||||
}[backbone]
|
||||
|
||||
if "next_vit" in backbone:
|
||||
in_features = {
|
||||
"next_vit_large_6m": [96, 256, 512, 1024],
|
||||
}[backbone]
|
||||
else:
|
||||
in_features = None
|
||||
|
||||
# Instantiate backbone and reassemble blocks
|
||||
self.pretrained, self.scratch = _make_encoder(
|
||||
backbone,
|
||||
features,
|
||||
False, # Set to true of you want to train from scratch, uses ImageNet weights
|
||||
groups=1,
|
||||
expand=False,
|
||||
exportable=False,
|
||||
hooks=hooks,
|
||||
use_readout=readout,
|
||||
in_features=in_features,
|
||||
)
|
||||
|
||||
self.number_layers = len(hooks) if hooks is not None else 4
|
||||
size_refinenet3 = None
|
||||
self.scratch.stem_transpose = None
|
||||
|
||||
if "beit" in backbone:
|
||||
self.forward_transformer = forward_beit
|
||||
elif "swin" in backbone:
|
||||
self.forward_transformer = forward_swin
|
||||
elif "next_vit" in backbone:
|
||||
from .backbones.next_vit import forward_next_vit
|
||||
self.forward_transformer = forward_next_vit
|
||||
elif "levit" in backbone:
|
||||
self.forward_transformer = forward_levit
|
||||
size_refinenet3 = 7
|
||||
self.scratch.stem_transpose = stem_b4_transpose(256, 128, get_act_layer("hard_swish"))
|
||||
else:
|
||||
self.forward_transformer = forward_vit
|
||||
|
||||
self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet3 = _make_fusion_block(features, use_bn, size_refinenet3)
|
||||
if self.number_layers >= 4:
|
||||
self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
|
||||
|
||||
self.scratch.output_conv = head
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
if self.channels_last == True:
|
||||
x.contiguous(memory_format=torch.channels_last)
|
||||
|
||||
layers = self.forward_transformer(self.pretrained, x)
|
||||
if self.number_layers == 3:
|
||||
layer_1, layer_2, layer_3 = layers
|
||||
else:
|
||||
layer_1, layer_2, layer_3, layer_4 = layers
|
||||
|
||||
layer_1_rn = self.scratch.layer1_rn(layer_1)
|
||||
layer_2_rn = self.scratch.layer2_rn(layer_2)
|
||||
layer_3_rn = self.scratch.layer3_rn(layer_3)
|
||||
if self.number_layers >= 4:
|
||||
layer_4_rn = self.scratch.layer4_rn(layer_4)
|
||||
|
||||
if self.number_layers == 3:
|
||||
path_3 = self.scratch.refinenet3(layer_3_rn, size=layer_2_rn.shape[2:])
|
||||
else:
|
||||
path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
|
||||
path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])
|
||||
path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])
|
||||
path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
|
||||
|
||||
if self.scratch.stem_transpose is not None:
|
||||
path_1 = self.scratch.stem_transpose(path_1)
|
||||
|
||||
out = self.scratch.output_conv(path_1)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class DPTDepthModel(DPT):
|
||||
def __init__(self, path=None, non_negative=True, **kwargs):
|
||||
features = kwargs["features"] if "features" in kwargs else 256
|
||||
head_features_1 = kwargs["head_features_1"] if "head_features_1" in kwargs else features
|
||||
head_features_2 = kwargs["head_features_2"] if "head_features_2" in kwargs else 32
|
||||
kwargs.pop("head_features_1", None)
|
||||
kwargs.pop("head_features_2", None)
|
||||
|
||||
head = nn.Sequential(
|
||||
nn.Conv2d(head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1),
|
||||
Interpolate(scale_factor=2, mode="bilinear", align_corners=True),
|
||||
nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),
|
||||
nn.ReLU(True) if non_negative else nn.Identity(),
|
||||
nn.Identity(),
|
||||
)
|
||||
|
||||
super().__init__(head, **kwargs)
|
||||
|
||||
if path is not None:
|
||||
self.load(path)
|
||||
|
||||
def forward(self, x):
|
||||
return super().forward(x).squeeze(dim=1)
|
||||
Reference in New Issue
Block a user