| | import torch |
| | import torch.nn as nn |
| | from torch import Tensor |
| |
|
| | from .base_model import BaseModel |
| | from .blocks import ( |
| | FeatureFusionBlock_custom, |
| | Interpolate, |
| | _make_encoder, |
| | 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, |
| | enable_attention_hooks=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, |
| | enable_attention_hooks=enable_attention_hooks, |
| | ) |
| |
|
| | 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: Tensor) -> Tensor: |
| | if self.channels_last == 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, scale=1.0, shift=0.0, invert=False, **kwargs |
| | ): |
| | features = kwargs["features"] if "features" in kwargs else 256 |
| |
|
| | self.scale = scale |
| | self.shift = shift |
| | self.invert = invert |
| |
|
| | 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: Tensor) -> Tensor: |
| | """Input x of shape [b, c, h, w] |
| | Return tensor of shape [b, c, h, w] |
| | """ |
| | inv_depth = super().forward(x) |
| |
|
| | if self.invert: |
| | depth = self.scale * inv_depth + self.shift |
| | depth[depth < 1e-8] = 1e-8 |
| | depth = 1.0 / depth |
| | return depth |
| | else: |
| | return inv_depth |
| |
|
| |
|
