| """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, Interpolate, _make_encoder |
|
|
|
|
| class MidasNet(BaseModel): |
| """Network for monocular depth estimation. |
| """ |
|
|
| def __init__(self, path=None, features=256, non_negative=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, self).__init__() |
|
|
| use_pretrained = False if path is None else True |
|
|
| self.pretrained, self.scratch = _make_encoder(backbone="resnext101_wsl", features=features, use_pretrained=use_pretrained) |
|
|
| self.scratch.refinenet4 = FeatureFusionBlock(features) |
| self.scratch.refinenet3 = FeatureFusionBlock(features) |
| self.scratch.refinenet2 = FeatureFusionBlock(features) |
| self.scratch.refinenet1 = FeatureFusionBlock(features) |
|
|
| self.scratch.output_conv = nn.Sequential( |
| nn.Conv2d(features, 128, kernel_size=3, stride=1, padding=1), |
| Interpolate(scale_factor=2, mode="bilinear"), |
| nn.Conv2d(128, 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(), |
| ) |
|
|
| if path: |
| self.load(path) |
|
|
| def forward(self, x): |
| """Forward pass. |
| |
| Args: |
| x (tensor): input data (image) |
| |
| Returns: |
| tensor: depth |
| """ |
|
|
| 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) |
|
|
