| | """MidashNet: Network for monocular depth estimation trained by mixing several datasets. |
| | This file contains code that is adapted from |
| | https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py |
| | """ |
| | import torch |
| | import torch.nn as nn |
| |
|
| | from .base_model import BaseModel |
| | from .blocks import FeatureFusionBlock, FeatureFusionBlock_custom, Interpolate, _make_encoder |
| |
|
| |
|
| | class MidasNet_small(BaseModel): |
| | """Network for monocular depth estimation. |
| | """ |
| |
|
| | def __init__(self, path=None, features=64, backbone="efficientnet_lite3", non_negative=True, exportable=True, channels_last=False, align_corners=True, |
| | blocks={'expand': True}): |
| | """Init. |
| | |
| | Args: |
| | path (str, optional): Path to saved model. Defaults to None. |
| | features (int, optional): Number of features. Defaults to 256. |
| | backbone (str, optional): Backbone network for encoder. Defaults to resnet50 |
| | """ |
| | print("Loading weights: ", path) |
| |
|
| | super(MidasNet_small, self).__init__() |
| |
|
| | use_pretrained = False if path else True |
| | |
| | self.channels_last = channels_last |
| | self.blocks = blocks |
| | self.backbone = backbone |
| |
|
| | self.groups = 1 |
| |
|
| | features1=features |
| | features2=features |
| | features3=features |
| | features4=features |
| | self.expand = False |
| | if "expand" in self.blocks and self.blocks['expand'] == True: |
| | self.expand = True |
| | features1=features |
| | features2=features*2 |
| | features3=features*4 |
| | features4=features*8 |
| |
|
| | self.pretrained, self.scratch = _make_encoder(self.backbone, features, use_pretrained, groups=self.groups, expand=self.expand, exportable=exportable) |
| | |
| | self.scratch.activation = nn.ReLU(False) |
| |
|
| | self.scratch.refinenet4 = FeatureFusionBlock_custom(features4, self.scratch.activation, deconv=False, bn=False, expand=self.expand, align_corners=align_corners) |
| | self.scratch.refinenet3 = FeatureFusionBlock_custom(features3, self.scratch.activation, deconv=False, bn=False, expand=self.expand, align_corners=align_corners) |
| | self.scratch.refinenet2 = FeatureFusionBlock_custom(features2, self.scratch.activation, deconv=False, bn=False, expand=self.expand, align_corners=align_corners) |
| | self.scratch.refinenet1 = FeatureFusionBlock_custom(features1, self.scratch.activation, deconv=False, bn=False, align_corners=align_corners) |
| |
|
| | |
| | self.scratch.output_conv = nn.Sequential( |
| | nn.Conv2d(features, features//2, kernel_size=3, stride=1, padding=1, groups=self.groups), |
| | Interpolate(scale_factor=2, mode="bilinear"), |
| | nn.Conv2d(features//2, 32, kernel_size=3, stride=1, padding=1), |
| | self.scratch.activation, |
| | nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0), |
| | nn.ReLU(True) if non_negative else nn.Identity(), |
| | nn.Identity(), |
| | ) |
| | |
| | if path: |
| | self.load(path) |
| |
|
| |
|
| | def forward(self, x): |
| | """Forward pass. |
| | |
| | Args: |
| | x (tensor): input data (image) |
| | |
| | Returns: |
| | tensor: depth |
| | """ |
| | if self.channels_last==True: |
| | print("self.channels_last = ", self.channels_last) |
| | x.contiguous(memory_format=torch.channels_last) |
| |
|
| |
|
| | layer_1 = self.pretrained.layer1(x) |
| | layer_2 = self.pretrained.layer2(layer_1) |
| | layer_3 = self.pretrained.layer3(layer_2) |
| | layer_4 = self.pretrained.layer4(layer_3) |
| | |
| | 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 torch.squeeze(out, dim=1) |
| |
|
| |
|
| |
|
| | def fuse_model(m): |
| | prev_previous_type = nn.Identity() |
| | prev_previous_name = '' |
| | previous_type = nn.Identity() |
| | previous_name = '' |
| | for name, module in m.named_modules(): |
| | if prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d and type(module) == nn.ReLU: |
| | |
| | torch.quantization.fuse_modules(m, [prev_previous_name, previous_name, name], inplace=True) |
| | elif prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d: |
| | |
| | torch.quantization.fuse_modules(m, [prev_previous_name, previous_name], inplace=True) |
| | |
| | |
| | |
| |
|
| | prev_previous_type = previous_type |
| | prev_previous_name = previous_name |
| | previous_type = type(module) |
| | previous_name = name |
