mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-12-23 20:43:19 +00:00
7c2f79fa98
* [pre-commit.ci] pre-commit autoupdate updates: - [github.com/PyCQA/autoflake: v2.2.1 → v2.3.1](https://github.com/PyCQA/autoflake/compare/v2.2.1...v2.3.1) - [github.com/pycqa/isort: 5.12.0 → 5.13.2](https://github.com/pycqa/isort/compare/5.12.0...5.13.2) - [github.com/psf/black-pre-commit-mirror: 23.9.1 → 24.4.2](https://github.com/psf/black-pre-commit-mirror/compare/23.9.1...24.4.2) - [github.com/pre-commit/mirrors-clang-format: v13.0.1 → v18.1.7](https://github.com/pre-commit/mirrors-clang-format/compare/v13.0.1...v18.1.7) - [github.com/pre-commit/pre-commit-hooks: v4.3.0 → v4.6.0](https://github.com/pre-commit/pre-commit-hooks/compare/v4.3.0...v4.6.0) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
"""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)
|