| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from .base_model import BaseModel |
| from .blocks import ( |
| FeatureFusionBlock, |
| FeatureFusionBlock_custom, |
| Interpolate, |
| _make_encoder, |
| forward_vit, |
| ) |
| from ..multi_task_head import MultitaskHead |
|
|
|
|
| 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, |
| use_layer_scale=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, |
| use_layer_scale=use_layer_scale, |
| ) |
|
|
| 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 == 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 DPTFieldModel(DPT): |
| def __init__(self, path=None, non_negative=True, head_size=[[3],[1],[1],[2],[2]], **kwargs): |
| features = kwargs["features"] if "features" in kwargs else 256 |
|
|
| kwargs["use_bn"] = True |
|
|
| num_class = sum(sum(head_size,[])) |
| head = nn.Sequential( |
| nn.Conv2d(features, features//2, kernel_size=3, stride=1, padding=1), |
| |
| nn.ReLU(True), |
| MultitaskHead(features//2, num_class, head_size=head_size), |
| ) |
|
|
| super().__init__(head, **kwargs) |
|
|
| self.stride = 2 |
|
|
| def forward(self, x): |
| if x.shape[1] == 1: |
| x = torch.cat([x,x,x], dim=1) |
| |
| out = super().forward(x) |
| return out, None |
|
|
|
|
