Spaces:
Running
on
Zero
Running
on
Zero
| # -*- coding: utf-8 -*- | |
| """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_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'] is 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 is 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: | |
| # print("FUSED ", prev_previous_name, previous_name, name) | |
| torch.quantization.fuse_modules( | |
| m, [prev_previous_name, previous_name, name], inplace=True) | |
| elif prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d: | |
| # print("FUSED ", prev_previous_name, previous_name) | |
| torch.quantization.fuse_modules( | |
| m, [prev_previous_name, previous_name], inplace=True) | |
| # elif previous_type == nn.Conv2d and type(module) == nn.ReLU: | |
| # print("FUSED ", previous_name, name) | |
| # torch.quantization.fuse_modules(m, [previous_name, name], inplace=True) | |
| prev_previous_type = previous_type | |
| prev_previous_name = previous_name | |
| previous_type = type(module) | |
| previous_name = name | |