| |
|
| |
|
| | import torch
|
| | import torch.nn as nn
|
| |
|
| | from .base_model import BaseModel
|
| | from .blocks import FeatureFusionBlock_custom, Interpolate, _make_encoder
|
| | from .vit import forward_vit
|
| |
|
| |
|
| | def _make_fusion_block(features, use_bn):
|
| | return FeatureFusionBlock_custom(
|
| | features,
|
| | nn.ReLU(False),
|
| | deconv=False,
|
| | bn=use_bn,
|
| | expand=False,
|
| | align_corners=True,
|
| | )
|
| |
|
| |
|
| | class DPT(BaseModel):
|
| | def __init__(
|
| | self,
|
| | head,
|
| | features=256,
|
| | backbone='vitb_rn50_384',
|
| | readout='project',
|
| | channels_last=False,
|
| | use_bn=False,
|
| | ):
|
| |
|
| | super(DPT, self).__init__()
|
| |
|
| | self.channels_last = channels_last
|
| |
|
| | hooks = {
|
| | 'vitb_rn50_384': [0, 1, 8, 11],
|
| | 'vitb16_384': [2, 5, 8, 11],
|
| | 'vitl16_384': [5, 11, 17, 23],
|
| | }
|
| |
|
| |
|
| | self.pretrained, self.scratch = _make_encoder(
|
| | backbone,
|
| | features,
|
| | False,
|
| | groups=1,
|
| | expand=False,
|
| | exportable=False,
|
| | hooks=hooks[backbone],
|
| | use_readout=readout,
|
| | )
|
| |
|
| | 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)
|
| | self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
|
| |
|
| | self.scratch.output_conv = head
|
| |
|
| | def forward(self, x):
|
| | if self.channels_last is True:
|
| | x.contiguous(memory_format=torch.channels_last)
|
| |
|
| | layer_1, layer_2, layer_3, layer_4 = forward_vit(self.pretrained, x)
|
| |
|
| | 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)
|
| | layer_4_rn = self.scratch.layer4_rn(layer_4)
|
| |
|
| | path_4 = self.scratch.refinenet4(layer_4_rn)
|
| | path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
|
| | path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
|
| | path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
|
| |
|
| | 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 = 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(True),
|
| | nn.Conv2d(32, 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)
|
| |
|
