RGBD-SOD/dptdepth

1import torch2import torch.nn as nn3from torch import Tensor4 5from .base_model import BaseModel6from .blocks import (7    FeatureFusionBlock_custom,8    Interpolate,9    _make_encoder,10    forward_vit,11)12 13 14def _make_fusion_block(features, use_bn):15    return FeatureFusionBlock_custom(16        features,17        nn.ReLU(False),18        deconv=False,19        bn=use_bn,20        expand=False,21        align_corners=True,22    )23 24 25class DPT(BaseModel):26    def __init__(27        self,28        head,29        features=256,30        backbone="vitb_rn50_384",31        readout="project",32        channels_last=False,33        use_bn=False,34        enable_attention_hooks=False,35    ):36 37        super(DPT, self).__init__()38 39        self.channels_last = channels_last40 41        hooks = {42            "vitb_rn50_384": [0, 1, 8, 11],43            "vitb16_384": [2, 5, 8, 11],44            "vitl16_384": [5, 11, 17, 23],45        }46 47        # Instantiate backbone and reassemble blocks48        self.pretrained, self.scratch = _make_encoder(49            backbone,50            features,51            False,  # Set to true of you want to train from scratch, uses ImageNet weights52            groups=1,53            expand=False,54            exportable=False,55            hooks=hooks[backbone],56            use_readout=readout,57            enable_attention_hooks=enable_attention_hooks,58        )59 60        self.scratch.refinenet1 = _make_fusion_block(features, use_bn)61        self.scratch.refinenet2 = _make_fusion_block(features, use_bn)62        self.scratch.refinenet3 = _make_fusion_block(features, use_bn)63        self.scratch.refinenet4 = _make_fusion_block(features, use_bn)64 65        self.scratch.output_conv = head66 67    def forward(self, x: Tensor) -> Tensor:68        if self.channels_last == True:69            x.contiguous(memory_format=torch.channels_last)70 71        layer_1, layer_2, layer_3, layer_4 = forward_vit(self.pretrained, x)72 73        layer_1_rn = self.scratch.layer1_rn(layer_1)74        layer_2_rn = self.scratch.layer2_rn(layer_2)75        layer_3_rn = self.scratch.layer3_rn(layer_3)76        layer_4_rn = self.scratch.layer4_rn(layer_4)77 78        path_4 = self.scratch.refinenet4(layer_4_rn)79        path_3 = self.scratch.refinenet3(path_4, layer_3_rn)80        path_2 = self.scratch.refinenet2(path_3, layer_2_rn)81        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)82 83        out = self.scratch.output_conv(path_1)84 85        return out86 87 88class DPTDepthModel(DPT):89    def __init__(90        self, path=None, non_negative=True, scale=1.0, shift=0.0, invert=False, **kwargs91    ):92        features = kwargs["features"] if "features" in kwargs else 25693 94        self.scale = scale95        self.shift = shift96        self.invert = invert97 98        head = nn.Sequential(99            nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1),100            Interpolate(scale_factor=2, mode="bilinear", align_corners=True),101            nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),102            nn.ReLU(True),103            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),104            nn.ReLU(True) if non_negative else nn.Identity(),105            nn.Identity(),106        )107 108        super().__init__(head, **kwargs)109 110        if path is not None:111            self.load(path)112 113    def forward(self, x: Tensor) -> Tensor:114        """Input x of shape [b, c, h, w]115        Return tensor of shape [b, c, h, w]116        """117        inv_depth = super().forward(x)118 119        if self.invert:120            depth = self.scale * inv_depth + self.shift121            depth[depth < 1e-8] = 1e-8122            depth = 1.0 / depth123            return depth124        else:125            return inv_depth126 127