From 9667548d9e16681b07aa84ebda0ebc11c9768d69 Mon Sep 17 00:00:00 2001 From: "shengzhe.sz" Date: Wed, 11 Jan 2023 10:23:24 +0800 Subject: [PATCH] add panorama_depth_estimation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11366392 * add panorama_depth_estimation: pipeline, model, test * modelhub：https://modelscope.cn/models/damo/cv_unifuse_panorama-depth-estimation/summary --- .../test/images/panorama_depth_estimation.jpg | 3 + modelscope/metainfo.py | 2 + modelscope/models/cv/__init__.py | 13 +- .../cv/panorama_depth_estimation/__init__.py | 22 + .../networks/__init__.py | 3 + .../networks/equi.py | 133 ++++++ .../networks/layers.py | 235 ++++++++++ .../networks/mobilenet.py | 238 ++++++++++ .../networks/resnet.py | 424 ++++++++++++++++++ .../networks/unifuse.py | 219 +++++++++ .../networks/util.py | 111 +++++ .../unifuse_model.py | 89 ++++ modelscope/pipelines/builder.py | 9 +- modelscope/pipelines/cv/__init__.py | 1 + .../cv/panorama_depth_estimation_pipeline.py | 85 ++++ modelscope/utils/constant.py | 1 + .../test_panorama_depth_estimation.py | 34 ++ 17 files changed, 1613 insertions(+), 9 deletions(-) create mode 100644 data/test/images/panorama_depth_estimation.jpg create mode 100644 modelscope/models/cv/panorama_depth_estimation/__init__.py create mode 100644 modelscope/models/cv/panorama_depth_estimation/networks/__init__.py create mode 100644 modelscope/models/cv/panorama_depth_estimation/networks/equi.py create mode 100644 modelscope/models/cv/panorama_depth_estimation/networks/layers.py create mode 100644 modelscope/models/cv/panorama_depth_estimation/networks/mobilenet.py create mode 100644 modelscope/models/cv/panorama_depth_estimation/networks/resnet.py create mode 100644 modelscope/models/cv/panorama_depth_estimation/networks/unifuse.py create mode 100644 modelscope/models/cv/panorama_depth_estimation/networks/util.py create mode 100644 modelscope/models/cv/panorama_depth_estimation/unifuse_model.py create mode 100644 modelscope/pipelines/cv/panorama_depth_estimation_pipeline.py create mode 100644 tests/pipelines/test_panorama_depth_estimation.py diff --git a/data/test/images/panorama_depth_estimation.jpg b/data/test/images/panorama_depth_estimation.jpg new file mode 100644 index 00000000..1c519732 --- /dev/null +++ b/data/test/images/panorama_depth_estimation.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0df5f2de59df6b55d8ee5d414cc2f98d714e14b518c159d4085ad2ac65d36627 +size 137606 diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 71596262..5c124b10 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -40,6 +40,7 @@ class Models(object): vitadapter_semantic_segmentation = 'vitadapter-semantic-segmentation' text_driven_segmentation = 'text-driven-segmentation' newcrfs_depth_estimation = 'newcrfs-depth-estimation' + unifuse_depth_estimation = 'unifuse-depth-estimation' dro_resnet18_depth_estimation = 'dro-resnet18-depth-estimation' resnet50_bert = 'resnet50-bert' referring_video_object_segmentation = 'swinT-referring-video-object-segmentation' @@ -257,6 +258,7 @@ class Pipelines(object): image_semantic_segmentation = 'image-semantic-segmentation' image_depth_estimation = 'image-depth-estimation' video_depth_estimation = 'video-depth-estimation' + panorama_depth_estimation = 'panorama-depth-estimation' image_reid_person = 'passvitb-image-reid-person' image_inpainting = 'fft-inpainting' text_driven_segmentation = 'text-driven-segmentation' diff --git a/modelscope/models/cv/__init__.py b/modelscope/models/cv/__init__.py index e5fa9818..50f37cb9 100644 --- a/modelscope/models/cv/__init__.py +++ b/modelscope/models/cv/__init__.py @@ -12,11 +12,12 @@ from . import (action_recognition, animal_recognition, body_2d_keypoints, image_semantic_segmentation, image_to_image_generation, image_to_image_translation, language_guided_video_summarization, movie_scene_segmentation, object_detection, - pointcloud_sceneflow_estimation, product_retrieval_embedding, - realtime_object_detection, referring_video_object_segmentation, - salient_detection, shop_segmentation, super_resolution, - video_frame_interpolation, video_object_segmentation, - video_single_object_tracking, video_stabilization, - video_summarization, video_super_resolution, virual_tryon) + panorama_depth_estimation, pointcloud_sceneflow_estimation, + product_retrieval_embedding, realtime_object_detection, + referring_video_object_segmentation, salient_detection, + shop_segmentation, super_resolution, video_frame_interpolation, + video_object_segmentation, video_single_object_tracking, + video_stabilization, video_summarization, + video_super_resolution, virual_tryon) # yapf: enable diff --git a/modelscope/models/cv/panorama_depth_estimation/__init__.py b/modelscope/models/cv/panorama_depth_estimation/__init__.py new file mode 100644 index 00000000..b20f6209 --- /dev/null +++ b/modelscope/models/cv/panorama_depth_estimation/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import TYPE_CHECKING + +from modelscope.utils.import_utils import LazyImportModule + +if TYPE_CHECKING: + from .unifuse_model import PanoramaDepthEstimation + +else: + _import_structure = { + 'unifuse_model': ['PanoramaDepthEstimation'], + } + + import sys + + sys.modules[__name__] = LazyImportModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/modelscope/models/cv/panorama_depth_estimation/networks/__init__.py b/modelscope/models/cv/panorama_depth_estimation/networks/__init__.py new file mode 100644 index 00000000..5f0fbd21 --- /dev/null +++ b/modelscope/models/cv/panorama_depth_estimation/networks/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from .equi import Equi +from .unifuse import UniFuse diff --git a/modelscope/models/cv/panorama_depth_estimation/networks/equi.py b/modelscope/models/cv/panorama_depth_estimation/networks/equi.py new file mode 100644 index 00000000..3b8bb3d3 --- /dev/null +++ b/modelscope/models/cv/panorama_depth_estimation/networks/equi.py @@ -0,0 +1,133 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from __future__ import absolute_import, division, print_function +from collections import OrderedDict + +import numpy as np +import torch +import torch.nn as nn + +from .layers import Conv3x3, ConvBlock, upsample +from .mobilenet import mobilenet_v2 +from .resnet import resnet18, resnet34, resnet50, resnet101, resnet152 + + +class Equi(nn.Module): + """ Model: Resnet based Encoder + Decoder + """ + + def __init__(self, + num_layers, + equi_h, + equi_w, + pretrained=False, + max_depth=10.0, + **kwargs): + super(Equi, self).__init__() + + self.num_layers = num_layers + self.equi_h = equi_h + self.equi_w = equi_w + self.cube_h = equi_h // 2 + + # encoder + encoder = { + 2: mobilenet_v2, + 18: resnet18, + 34: resnet34, + 50: resnet50, + 101: resnet101, + 152: resnet152 + } + + if num_layers not in encoder: + raise ValueError( + '{} is not a valid number of resnet layers'.format(num_layers)) + self.equi_encoder = encoder[num_layers](pretrained) + + self.num_ch_enc = np.array([64, 64, 128, 256, 512]) + if num_layers > 34: + self.num_ch_enc[1:] *= 4 + if num_layers < 18: + self.num_ch_enc = np.array([16, 24, 32, 96, 320]) + + # decoder + self.num_ch_dec = np.array([16, 32, 64, 128, 256]) + self.equi_dec_convs = OrderedDict() + + self.equi_dec_convs['upconv_5'] = ConvBlock(self.num_ch_enc[4], + self.num_ch_dec[4]) + + self.equi_dec_convs['deconv_4'] = ConvBlock( + self.num_ch_dec[4] + self.num_ch_enc[3], self.num_ch_dec[4]) + self.equi_dec_convs['upconv_4'] = ConvBlock(self.num_ch_dec[4], + self.num_ch_dec[3]) + + self.equi_dec_convs['deconv_3'] = ConvBlock( + self.num_ch_dec[3] + self.num_ch_enc[2], self.num_ch_dec[3]) + self.equi_dec_convs['upconv_3'] = ConvBlock(self.num_ch_dec[3], + self.num_ch_dec[2]) + + self.equi_dec_convs['deconv_2'] = ConvBlock( + self.num_ch_dec[2] + self.num_ch_enc[1], self.num_ch_dec[2]) + self.equi_dec_convs['upconv_2'] = ConvBlock(self.num_ch_dec[2], + self.num_ch_dec[1]) + + self.equi_dec_convs['deconv_1'] = ConvBlock( + self.num_ch_dec[1] + self.num_ch_enc[0], self.num_ch_dec[1]) + self.equi_dec_convs['upconv_1'] = ConvBlock(self.num_ch_dec[1], + self.num_ch_dec[0]) + + self.equi_dec_convs['deconv_0'] = ConvBlock(self.num_ch_dec[0], + self.num_ch_dec[0]) + self.equi_dec_convs['depthconv_0'] = Conv3x3(self.num_ch_dec[0], 1) + + self.equi_decoder = nn.ModuleList(list(self.equi_dec_convs.values())) + + self.sigmoid = nn.Sigmoid() + self.max_depth = nn.Parameter( + torch.tensor(max_depth), requires_grad=False) + + def forward(self, input_equi_image, input_cube_image): + + # euqi image encoding + if self.num_layers < 18: + equi_enc_feat0, equi_enc_feat1, equi_enc_feat2, equi_enc_feat3, equi_enc_feat4 \ + = self.equi_encoder(input_equi_image) + else: + x = self.equi_encoder.conv1(input_equi_image) + x = self.equi_encoder.relu(self.equi_encoder.bn1(x)) + equi_enc_feat0 = x + + x = self.equi_encoder.maxpool(x) + equi_enc_feat1 = self.equi_encoder.layer1(x) + equi_enc_feat2 = self.equi_encoder.layer2(equi_enc_feat1) + equi_enc_feat3 = self.equi_encoder.layer3(equi_enc_feat2) + equi_enc_feat4 = self.equi_encoder.layer4(equi_enc_feat3) + + # euqi image decoding + outputs = {} + + equi_x = equi_enc_feat4 + equi_x = upsample(self.equi_dec_convs['upconv_5'](equi_x)) + + equi_x = torch.cat([equi_x, equi_enc_feat3], 1) + equi_x = self.equi_dec_convs['deconv_4'](equi_x) + equi_x = upsample(self.equi_dec_convs['upconv_4'](equi_x)) + + equi_x = torch.cat([equi_x, equi_enc_feat2], 1) + equi_x = self.equi_dec_convs['deconv_3'](equi_x) + equi_x = upsample(self.equi_dec_convs['upconv_3'](equi_x)) + + equi_x = torch.cat([equi_x, equi_enc_feat1], 1) + equi_x = self.equi_dec_convs['deconv_2'](equi_x) + equi_x = upsample(self.equi_dec_convs['upconv_2'](equi_x)) + + equi_x = torch.cat([equi_x, equi_enc_feat0], 1) + equi_x = self.equi_dec_convs['deconv_1'](equi_x) + equi_x = upsample(self.equi_dec_convs['upconv_1'](equi_x)) + + equi_x = self.equi_dec_convs['deconv_0'](equi_x) + equi_depth = self.equi_dec_convs['depthconv_0'](equi_x) + outputs['pred_depth'] = self.max_depth * self.sigmoid(equi_depth) + + return outputs diff --git a/modelscope/models/cv/panorama_depth_estimation/networks/layers.py b/modelscope/models/cv/panorama_depth_estimation/networks/layers.py new file mode 100644 index 00000000..99e166aa --- /dev/null +++ b/modelscope/models/cv/panorama_depth_estimation/networks/layers.py @@ -0,0 +1,235 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class Conv3x3(nn.Module): + """Layer to pad and convolve input + """ + + def __init__(self, in_channels, out_channels, bias=True): + super(Conv3x3, self).__init__() + + self.pad = nn.ZeroPad2d(1) + self.conv = nn.Conv2d( + int(in_channels), int(out_channels), 3, bias=bias) + + def forward(self, x): + out = self.pad(x) + out = self.conv(out) + return out + + +class ConvBlock(nn.Module): + """Layer to perform a convolution followed by ELU + """ + + def __init__(self, in_channels, out_channels, bias=True): + super(ConvBlock, self).__init__() + + self.conv = Conv3x3(in_channels, out_channels, bias) + self.nonlin = nn.ELU(inplace=True) + + def forward(self, x): + out = self.conv(x) + out = self.nonlin(out) + return out + + +def upsample(x): + """Upsample input tensor by a factor of 2 + """ + return F.interpolate(x, scale_factor=2, mode='nearest') + + +# Based on https://github.com/sunset1995/py360convert +class Cube2Equirec(nn.Module): + + def __init__(self, face_w, equ_h, equ_w): + super(Cube2Equirec, self).__init__() + ''' + face_w: int, the length of each face of the cubemap + equ_h: int, height of the equirectangular image + equ_w: int, width of the equirectangular image + ''' + + self.face_w = face_w + self.equ_h = equ_h + self.equ_w = equ_w + + # Get face id to each pixel: 0F 1R 2B 3L 4U 5D + self._equirect_facetype() + self._equirect_faceuv() + + def _equirect_facetype(self): + ''' + 0F 1R 2B 3L 4U 5D + ''' + tp = np.roll( + np.arange(4).repeat(self.equ_w // 4)[None, :].repeat( + self.equ_h, 0), 3 * self.equ_w // 8, 1) + + # Prepare ceil mask + mask = np.zeros((self.equ_h, self.equ_w // 4), np.bool) + idx = np.linspace(-np.pi, np.pi, self.equ_w // 4) / 4 + idx = self.equ_h // 2 - np.round( + np.arctan(np.cos(idx)) * self.equ_h / np.pi).astype(int) + for i, j in enumerate(idx): + mask[:j, i] = 1 + mask = np.roll(np.concatenate([mask] * 4, 1), 3 * self.equ_w // 8, 1) + + tp[mask] = 4 + tp[np.flip(mask, 0)] = 5 + + self.tp = tp + self.mask = mask + + def _equirect_faceuv(self): + + lon = ( + (np.linspace(0, self.equ_w - 1, num=self.equ_w, dtype=np.float32) + + 0.5) / self.equ_w - 0.5) * 2 * np.pi + lat = -( + (np.linspace(0, self.equ_h - 1, num=self.equ_h, dtype=np.float32) + + 0.5) / self.equ_h - 0.5) * np.pi + + lon, lat = np.meshgrid(lon, lat) + + coor_u = np.zeros((self.equ_h, self.equ_w), dtype=np.float32) + coor_v = np.zeros((self.equ_h, self.equ_w), dtype=np.float32) + + for i in range(4): + mask = (self.tp == i) + coor_u[mask] = 0.5 * np.tan(lon[mask] - np.pi * i / 2) + coor_v[mask] = -0.5 * np.tan( + lat[mask]) / np.cos(lon[mask] - np.pi * i / 2) + + mask = (self.tp == 4) + c = 0.5 * np.tan(np.pi / 2 - lat[mask]) + coor_u[mask] = c * np.sin(lon[mask]) + coor_v[mask] = c * np.cos(lon[mask]) + + mask = (self.tp == 5) + c = 0.5 * np.tan(np.pi / 2 - np.abs(lat[mask])) + coor_u[mask] = c * np.sin(lon[mask]) + coor_v[mask] = -c * np.cos(lon[mask]) + + # Final renormalize + coor_u = (np.clip(coor_u, -0.5, 0.5)) * 2 + coor_v = (np.clip(coor_v, -0.5, 0.5)) * 2 + + # Convert to torch tensor + self.tp = torch.from_numpy(self.tp.astype(np.float32) / 2.5 - 1) + self.coor_u = torch.from_numpy(coor_u) + self.coor_v = torch.from_numpy(coor_v) + + sample_grid = torch.stack([self.coor_u, self.coor_v, self.tp], + dim=-1).view(1, 1, self.equ_h, self.equ_w, 3) + self.sample_grid = nn.Parameter(sample_grid, requires_grad=False) + + def forward(self, cube_feat): + + bs, ch, h, w = cube_feat.shape + assert h == self.face_w and w // 6 == self.face_w + + cube_feat = cube_feat.view(bs, ch, 1, h, w) + cube_feat = torch.cat( + torch.split(cube_feat, self.face_w, dim=-1), dim=2) + + cube_feat = cube_feat.view([bs, ch, 6, self.face_w, self.face_w]) + sample_grid = torch.cat(bs * [self.sample_grid], dim=0) + equi_feat = F.grid_sample( + cube_feat, sample_grid, padding_mode='border', align_corners=True) + + return equi_feat.squeeze(2) + + +class Concat(nn.Module): + + def __init__(self, channels, **kwargs): + super(Concat, self).__init__() + self.conv = nn.Conv2d(channels * 2, channels, 1, bias=False) + self.relu = nn.ReLU(inplace=True) + + def forward(self, equi_feat, c2e_feat): + + x = torch.cat([equi_feat, c2e_feat], 1) + x = self.relu(self.conv(x)) + return x + + +# Based on https://github.com/Yeh-yu-hsuan/BiFuse/blob/master/models/FCRN.py +class BiProj(nn.Module): + + def __init__(self, channels, **kwargs): + super(BiProj, self).__init__() + + self.conv_c2e = nn.Sequential( + nn.Conv2d(channels, channels, kernel_size=3, padding=1), + nn.ReLU(inplace=True)) + self.conv_e2c = nn.Sequential( + nn.Conv2d(channels, channels, kernel_size=3, padding=1), + nn.ReLU(inplace=True)) + self.conv_mask = nn.Sequential( + nn.Conv2d(channels * 2, 1, kernel_size=1, padding=0), nn.Sigmoid()) + + def forward(self, equi_feat, c2e_feat): + aaa = self.conv_e2c(equi_feat) + tmp_equi = self.conv_c2e(c2e_feat) + mask_equi = self.conv_mask(torch.cat([aaa, tmp_equi], dim=1)) + tmp_equi = tmp_equi.clone() * mask_equi + return equi_feat + tmp_equi + + +# from https://github.com/moskomule/senet.pytorch/blob/master/senet/se_module.py +class SELayer(nn.Module): + + def __init__(self, channel, reduction=16): + super(SELayer, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channel, channel // reduction, bias=False), + nn.ReLU(inplace=True), + nn.Linear(channel // reduction, channel, bias=False), nn.Sigmoid()) + + def forward(self, x): + b, c, _, _ = x.size() + y = self.avg_pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y.expand_as(x) + + +class CEELayer(nn.Module): + + def __init__(self, channels, SE=True): + super(CEELayer, self).__init__() + + self.res_conv1 = nn.Conv2d( + channels * 2, channels, kernel_size=1, padding=0, bias=False) + self.res_bn1 = nn.BatchNorm2d(channels) + + self.res_conv2 = nn.Conv2d( + channels, channels, kernel_size=3, padding=1, bias=False) + self.res_bn2 = nn.BatchNorm2d(channels) + self.relu = nn.ReLU(inplace=True) + + self.SE = SE + if self.SE: + self.selayer = SELayer(channels * 2) + + self.conv = nn.Conv2d(channels * 2, channels, 1, bias=False) + + def forward(self, equi_feat, c2e_feat): + + x = torch.cat([equi_feat, c2e_feat], 1) + x = self.relu(self.res_bn1(self.res_conv1(x))) + shortcut = self.res_bn2(self.res_conv2(x)) + + x = c2e_feat + shortcut + x = torch.cat([equi_feat, x], 1) + if self.SE: + x = self.selayer(x) + x = self.relu(self.conv(x)) + return x diff --git a/modelscope/models/cv/panorama_depth_estimation/networks/mobilenet.py b/modelscope/models/cv/panorama_depth_estimation/networks/mobilenet.py new file mode 100644 index 00000000..d49b9a05 --- /dev/null +++ b/modelscope/models/cv/panorama_depth_estimation/networks/mobilenet.py @@ -0,0 +1,238 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +# Modified from https://github.com/pytorch/vision/blob/master/torchvision/models/mobilenet.py +from torch import nn + +try: + from torch.hub import load_state_dict_from_url +except ImportError: + from torch.utils.model_zoo import load_url as load_state_dict_from_url + +__all__ = ['MobileNetV2', 'mobilenet_v2'] + + +def _make_divisible(v, divisor, min_value=None): + """ + This function is taken from the original tf repo. + It ensures that all layers have a channel number that is divisible by 8 + It can be seen here: + https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py + :param v: + :param divisor: + :param min_value: + :return: + """ + if min_value is None: + min_value = divisor + new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) + # Make sure that round down does not go down by more than 10%. + if new_v < 0.9 * v: + new_v += divisor + return new_v + + +class ConvBNReLU(nn.Sequential): + + def __init__(self, + in_planes, + out_planes, + kernel_size=3, + stride=1, + groups=1, + norm_layer=None): + padding = (kernel_size - 1) // 2 + if norm_layer is None: + norm_layer = nn.BatchNorm2d + super(ConvBNReLU, self).__init__( + nn.Conv2d( + in_planes, + out_planes, + kernel_size, + stride, + padding, + groups=groups, + bias=False), norm_layer(out_planes), nn.ReLU6(inplace=True)) + + +class InvertedResidual(nn.Module): + + def __init__(self, inp, oup, stride, expand_ratio, norm_layer=None): + super(InvertedResidual, self).__init__() + self.stride = stride + assert stride in [1, 2] + + if norm_layer is None: + norm_layer = nn.BatchNorm2d + + hidden_dim = int(round(inp * expand_ratio)) + self.use_res_connect = self.stride == 1 and inp == oup + + layers = [] + if expand_ratio != 1: + # pw + layers.append( + ConvBNReLU( + inp, hidden_dim, kernel_size=1, norm_layer=norm_layer)) + layers.extend([ + # dw + ConvBNReLU( + hidden_dim, + hidden_dim, + stride=stride, + groups=hidden_dim, + norm_layer=norm_layer), + # pw-linear + nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), + norm_layer(oup), + ]) + self.conv = nn.Sequential(*layers) + + def forward(self, x): + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + +class MobileNetV2(nn.Module): + + def __init__(self, + width_mult=1.0, + inverted_residual_setting=None, + round_nearest=8, + block=None, + norm_layer=None): + """ + MobileNet V2 main class + + Args: + num_classes (int): Number of classes + width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount + inverted_residual_setting: Network structure + round_nearest (int): Round the number of channels in each layer to be a multiple of this number + Set to 1 to turn off rounding + block: Module specifying inverted residual building block for mobilenet + norm_layer: Module specifying the normalization layer to use + + """ + super(MobileNetV2, self).__init__() + + if block is None: + block = InvertedResidual + + if norm_layer is None: + norm_layer = nn.BatchNorm2d + + input_channel = 32 + + if inverted_residual_setting is None: + inverted_residual_setting = [ + # t, c, n, s + [1, 16, 1, 1], + [6, 24, 2, 2], + [6, 32, 3, 2], + [6, 64, 4, 2], + [6, 96, 3, 1], + [6, 160, 3, 2], + [6, 320, 1, 1], + ] + + # only check the first element, assuming user knows t,c,n,s are required + if len(inverted_residual_setting) == 0 or len( + inverted_residual_setting[0]) != 4: + raise ValueError('inverted_residual_setting should be non-empty ' + 'or a 4-element list, got {}'.format( + inverted_residual_setting)) + + # building first layer + input_channel = _make_divisible(input_channel * width_mult, + round_nearest) + features = [ + ConvBNReLU(3, input_channel, stride=2, norm_layer=norm_layer) + ] + # building inverted residual blocks + for t, c, n, s in inverted_residual_setting: + output_channel = _make_divisible(c * width_mult, round_nearest) + for i in range(n): + stride = s if i == 0 else 1 + features.append( + block( + input_channel, + output_channel, + stride, + expand_ratio=t, + norm_layer=norm_layer)) + input_channel = output_channel + # building last several layers + # features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1, norm_layer=norm_layer)) + # make it nn.Sequential + self.features = nn.Sequential(*features) + """ + # remove fcn as we don't need it for a depth prediction task + # building classifier + self.classifier = nn.Sequential( + nn.Dropout(0.2), + nn.Linear(self.last_channel, num_classes), + ) + """ + + # weight initialization + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out') + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + nn.init.zeros_(m.bias) + + def _forward_impl(self, x): + # This exists since TorchScript doesn't support inheritance, so the superclass method + # (this one) needs to have a name other than `forward` that can be accessed in a subclass + # Cannot use "squeeze" as batch-size can be 1 => must use reshape with x.shape[0] + + st = 0 + for i in range(2): + x = self.features[st + i](x) + st = st + 2 + feat0 = x + + for i in range(2): + x = self.features[st + i](x) + st = st + 2 + feat1 = x + + for i in range(3): + x = self.features[st + i](x) + st = st + 3 + feat2 = x + + for i in range(7): + x = self.features[st + i](x) + st = st + 7 + feat3 = x + + for i in range(4): + x = self.features[st + i](x) + feat4 = x + + return feat0, feat1, feat2, feat3, feat4 + + def forward(self, x): + return self._forward_impl(x) + + +def mobilenet_v2(pretrained=False, progress=True, **kwargs): + """ + Constructs a MobileNetV2 architecture from + `"MobileNetV2: Inverted Residuals and Linear Bottlenecks" `_. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet (deprecated) + progress (bool): If True, displays a progress bar of the download to stderr + """ + model = MobileNetV2(**kwargs) + + return model diff --git a/modelscope/models/cv/panorama_depth_estimation/networks/resnet.py b/modelscope/models/cv/panorama_depth_estimation/networks/resnet.py new file mode 100644 index 00000000..e6e6b1d3 --- /dev/null +++ b/modelscope/models/cv/panorama_depth_estimation/networks/resnet.py @@ -0,0 +1,424 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +# Modified from https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py +import torch +import torch.nn as nn + +try: + from torch.hub import load_state_dict_from_url +except ImportError: + from torch.utils.model_zoo import load_url as load_state_dict_from_url + +__all__ = [ + 'ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', + 'resnext50_32x4d', 'resnext101_32x8d', 'wide_resnet50_2', + 'wide_resnet101_2' +] + + +def conv3x3(in_planes, out_planes, padding, stride=1, groups=1, dilation=1): + """3x3 convolution with padding""" + return nn.Conv2d( + in_planes, + out_planes, + kernel_size=3, + stride=stride, + padding=padding, + groups=groups, + bias=False, + dilation=dilation) + + +def conv1x1(in_planes, out_planes, stride=1): + """1x1 convolution""" + return nn.Conv2d( + in_planes, out_planes, kernel_size=1, stride=stride, bias=False) + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, + inplanes, + planes, + stride=1, + downsample=None, + groups=1, + base_width=64, + dilation=1, + norm_layer=None): + super(BasicBlock, self).__init__() + if norm_layer is None: + norm_layer = nn.BatchNorm2d + if groups != 1 or base_width != 64: + raise ValueError( + 'BasicBlock only supports groups=1 and base_width=64') + if dilation > 1: + raise NotImplementedError( + 'Dilation > 1 not supported in BasicBlock') + # Both self.conv1 and self.downsample layers downsample the input when stride != 1 + + self.conv1 = conv3x3(inplanes, planes, 1, stride) + + self.bn1 = norm_layer(planes) + self.relu = nn.ReLU(inplace=True) + + self.conv2 = conv3x3(planes, planes, 1) + + self.bn2 = norm_layer(planes) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + identity = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + identity = self.downsample(identity) + + out += identity + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2) + # while original implementation places the stride at the first 1x1 convolution(self.conv1) + # according to "Deep residual learning for image recognition" https://arxiv.org/abs/1512.03385. + # This variant is also known as ResNet V1.5 and improves accuracy according to + # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch. + + expansion = 4 + + def __init__(self, + inplanes, + planes, + stride=1, + downsample=None, + groups=1, + base_width=64, + dilation=1, + norm_layer=None): + super(Bottleneck, self).__init__() + if norm_layer is None: + norm_layer = nn.BatchNorm2d + width = int(planes * (base_width / 64.)) * groups + + # Both self.conv2 and self.downsample layers downsample the input when stride != 1 + self.conv1 = conv1x1(inplanes, width) + self.bn1 = norm_layer(width) + + self.conv2 = conv3x3(width, width, 1, stride, groups, dilation) + + self.bn2 = norm_layer(width) + self.conv3 = conv1x1(width, planes * self.expansion) + self.bn3 = norm_layer(planes * self.expansion) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + identity = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu(out) + + return out + + +class ResNet(nn.Module): + + def __init__(self, + block, + layers, + num_input_images=1, + zero_init_residual=False, + groups=1, + width_per_group=64, + replace_stride_with_dilation=None, + norm_layer=None): + super(ResNet, self).__init__() + + if norm_layer is None: + norm_layer = nn.BatchNorm2d + self._norm_layer = norm_layer + + self.inplanes = 64 + self.dilation = 1 + if replace_stride_with_dilation is None: + # each element in the tuple indicates if we should replace + # the 2x2 stride with a dilated convolution instead + replace_stride_with_dilation = [False, False, False] + if len(replace_stride_with_dilation) != 3: + raise ValueError('replace_stride_with_dilation should be None ' + 'or a 3-element tuple, got {}'.format( + replace_stride_with_dilation)) + self.groups = groups + self.base_width = width_per_group + + self.conv1 = nn.Conv2d( + 3 * num_input_images, + self.inplanes, + kernel_size=7, + stride=2, + padding=3, + bias=False, + padding_mode='zeros') + + self.bn1 = norm_layer(self.inplanes) + self.relu = nn.ReLU(inplace=True) + + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer( + block, + 128, + layers[1], + stride=2, + dilate=replace_stride_with_dilation[0]) + self.layer3 = self._make_layer( + block, + 256, + layers[2], + stride=2, + dilate=replace_stride_with_dilation[1]) + self.layer4 = self._make_layer( + block, + 512, + layers[3], + stride=2, + dilate=replace_stride_with_dilation[2]) + """ + # remove fcn as we don't need it for a depth prediction task + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = nn.Linear(512 * block.expansion, num_classes) + """ + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_( + m.weight, mode='fan_out', nonlinearity='relu') + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + # Zero-initialize the last BN in each residual branch, + # so that the residual branch starts with zeros, and each residual block behaves like an identity. + # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 + if zero_init_residual: + for m in self.modules(): + if isinstance(m, Bottleneck): + nn.init.constant_(m.bn3.weight, 0) + elif isinstance(m, BasicBlock): + nn.init.constant_(m.bn2.weight, 0) + + def _make_layer(self, block, planes, blocks, stride=1, dilate=False): + norm_layer = self._norm_layer + downsample = None + previous_dilation = self.dilation + if dilate: + self.dilation *= stride + stride = 1 + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + conv1x1(self.inplanes, planes * block.expansion, stride), + norm_layer(planes * block.expansion), + ) + + layers = [] + layers.append( + block(self.inplanes, planes, stride, downsample, self.groups, + self.base_width, previous_dilation, norm_layer)) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append( + block( + self.inplanes, + planes, + groups=self.groups, + base_width=self.base_width, + dilation=self.dilation, + norm_layer=norm_layer)) + + return nn.Sequential(*layers) + + def _forward_impl(self, x): + # See note [TorchScript super()] + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + """ + x = self.avgpool(x) + x = torch.flatten(x, 1) + x = self.fc(x) + """ + + return x + + def forward(self, x): + return self._forward_impl(x) + + +def _resnet(arch, + block, + layers, + pretrained, + progress, + num_input_images=1, + **kwargs): + model = ResNet(block, layers, num_input_images=num_input_images, **kwargs) + + return model + + +def resnet18(pretrained=False, progress=True, **kwargs): + r"""ResNet-18 model from + `"Deep Residual Learning for Image Recognition" `_ + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet (deprecated) + progress (bool): If True, displays a progress bar of the download to stderr + """ + return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, + **kwargs) + + +def resnet34(pretrained=False, progress=True, **kwargs): + r"""ResNet-34 model from + `"Deep Residual Learning for Image Recognition" `_ + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, + **kwargs) + + +def resnet50(pretrained=False, progress=True, **kwargs): + r"""ResNet-50 model from + `"Deep Residual Learning for Image Recognition" `_ + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, + **kwargs) + + +def resnet101(pretrained=False, progress=True, **kwargs): + r"""ResNet-101 model from + `"Deep Residual Learning for Image Recognition" `_ + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, + progress, **kwargs) + + +def resnet152(pretrained=False, progress=True, **kwargs): + r"""ResNet-152 model from + `"Deep Residual Learning for Image Recognition" `_ + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, + progress, **kwargs) + + +def resnext50_32x4d(pretrained=False, progress=True, **kwargs): + r"""ResNeXt-50 32x4d model from + `"Aggregated Residual Transformation for Deep Neural Networks" `_ + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + kwargs['groups'] = 32 + kwargs['width_per_group'] = 4 + return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3], pretrained, + progress, **kwargs) + + +def resnext101_32x8d(pretrained=False, progress=True, **kwargs): + r"""ResNeXt-101 32x8d model from + `"Aggregated Residual Transformation for Deep Neural Networks" `_ + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + kwargs['groups'] = 32 + kwargs['width_per_group'] = 8 + return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3], pretrained, + progress, **kwargs) + + +def wide_resnet50_2(pretrained=False, progress=True, **kwargs): + r"""Wide ResNet-50-2 model from + `"Wide Residual Networks" `_ + + The model is the same as ResNet except for the bottleneck number of channels + which is twice larger in every block. The number of channels in outer 1x1 + convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 + channels, and in Wide ResNet-50-2 has 2048-1024-2048. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + kwargs['width_per_group'] = 64 * 2 + return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3], pretrained, + progress, **kwargs) + + +def wide_resnet101_2(pretrained=False, progress=True, **kwargs): + r"""Wide ResNet-101-2 model from + `"Wide Residual Networks" `_ + + The model is the same as ResNet except for the bottleneck number of channels + which is twice larger in every block. The number of channels in outer 1x1 + convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 + channels, and in Wide ResNet-50-2 has 2048-1024-2048. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + kwargs['width_per_group'] = 64 * 2 + return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3], pretrained, + progress, **kwargs) diff --git a/modelscope/models/cv/panorama_depth_estimation/networks/unifuse.py b/modelscope/models/cv/panorama_depth_estimation/networks/unifuse.py new file mode 100644 index 00000000..5715c89b --- /dev/null +++ b/modelscope/models/cv/panorama_depth_estimation/networks/unifuse.py @@ -0,0 +1,219 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from __future__ import absolute_import, division, print_function +from collections import OrderedDict + +import numpy as np +import torch +import torch.nn as nn + +from .layers import (BiProj, CEELayer, Concat, Conv3x3, ConvBlock, + Cube2Equirec, upsample) +from .mobilenet import mobilenet_v2 +from .resnet import resnet18, resnet34, resnet50, resnet101, resnet152 + + +class UniFuse(nn.Module): + """ UniFuse Model: Resnet based Euqi Encoder and Cube Encoder + Euqi Decoder + """ + + def __init__(self, + num_layers, + equi_h, + equi_w, + pretrained=False, + max_depth=10.0, + fusion_type='cee', + se_in_fusion=True): + super(UniFuse, self).__init__() + + self.num_layers = num_layers + self.equi_h = equi_h + self.equi_w = equi_w + self.cube_h = equi_h // 2 + + self.fusion_type = fusion_type + self.se_in_fusion = se_in_fusion + + # encoder + encoder = { + 2: mobilenet_v2, + 18: resnet18, + 34: resnet34, + 50: resnet50, + 101: resnet101, + 152: resnet152 + } + + if num_layers not in encoder: + raise ValueError( + '{} is not a valid number of resnet layers'.format(num_layers)) + self.equi_encoder = encoder[num_layers](pretrained) + self.cube_encoder = encoder[num_layers](pretrained) + + self.num_ch_enc = np.array([64, 64, 128, 256, 512]) + if num_layers > 34: + self.num_ch_enc[1:] *= 4 + + if num_layers < 18: + self.num_ch_enc = np.array([16, 24, 32, 96, 320]) + + # decoder + self.num_ch_dec = np.array([16, 32, 64, 128, 256]) + self.equi_dec_convs = OrderedDict() + self.c2e = {} + + Fusion_dict = {'cat': Concat, 'biproj': BiProj, 'cee': CEELayer} + FusionLayer = Fusion_dict[self.fusion_type] + + self.c2e['5'] = Cube2Equirec(self.cube_h // 32, self.equi_h // 32, + self.equi_w // 32) + + self.equi_dec_convs['fusion_5'] = FusionLayer( + self.num_ch_enc[4], SE=self.se_in_fusion) + self.equi_dec_convs['upconv_5'] = ConvBlock(self.num_ch_enc[4], + self.num_ch_dec[4]) + + self.c2e['4'] = Cube2Equirec(self.cube_h // 16, self.equi_h // 16, + self.equi_w // 16) + self.equi_dec_convs['fusion_4'] = FusionLayer( + self.num_ch_enc[3], SE=self.se_in_fusion) + self.equi_dec_convs['deconv_4'] = ConvBlock( + self.num_ch_dec[4] + self.num_ch_enc[3], self.num_ch_dec[4]) + self.equi_dec_convs['upconv_4'] = ConvBlock(self.num_ch_dec[4], + self.num_ch_dec[3]) + + self.c2e['3'] = Cube2Equirec(self.cube_h // 8, self.equi_h // 8, + self.equi_w // 8) + self.equi_dec_convs['fusion_3'] = FusionLayer( + self.num_ch_enc[2], SE=self.se_in_fusion) + self.equi_dec_convs['deconv_3'] = ConvBlock( + self.num_ch_dec[3] + self.num_ch_enc[2], self.num_ch_dec[3]) + self.equi_dec_convs['upconv_3'] = ConvBlock(self.num_ch_dec[3], + self.num_ch_dec[2]) + + self.c2e['2'] = Cube2Equirec(self.cube_h // 4, self.equi_h // 4, + self.equi_w // 4) + self.equi_dec_convs['fusion_2'] = FusionLayer( + self.num_ch_enc[1], SE=self.se_in_fusion) + self.equi_dec_convs['deconv_2'] = ConvBlock( + self.num_ch_dec[2] + self.num_ch_enc[1], self.num_ch_dec[2]) + self.equi_dec_convs['upconv_2'] = ConvBlock(self.num_ch_dec[2], + self.num_ch_dec[1]) + + self.c2e['1'] = Cube2Equirec(self.cube_h // 2, self.equi_h // 2, + self.equi_w // 2) + self.equi_dec_convs['fusion_1'] = FusionLayer( + self.num_ch_enc[0], SE=self.se_in_fusion) + self.equi_dec_convs['deconv_1'] = ConvBlock( + self.num_ch_dec[1] + self.num_ch_enc[0], self.num_ch_dec[1]) + self.equi_dec_convs['upconv_1'] = ConvBlock(self.num_ch_dec[1], + self.num_ch_dec[0]) + + self.equi_dec_convs['deconv_0'] = ConvBlock(self.num_ch_dec[0], + self.num_ch_dec[0]) + + self.equi_dec_convs['depthconv_0'] = Conv3x3(self.num_ch_dec[0], 1) + + self.equi_decoder = nn.ModuleList(list(self.equi_dec_convs.values())) + self.projectors = nn.ModuleList(list(self.c2e.values())) + + self.sigmoid = nn.Sigmoid() + + self.max_depth = nn.Parameter( + torch.tensor(max_depth), requires_grad=False) + + def forward(self, input_equi_image, input_cube_image): + + # euqi image encoding + + if self.num_layers < 18: + equi_enc_feat0, equi_enc_feat1, equi_enc_feat2, equi_enc_feat3, equi_enc_feat4 \ + = self.equi_encoder(input_equi_image) + else: + x = self.equi_encoder.conv1(input_equi_image) + x = self.equi_encoder.relu(self.equi_encoder.bn1(x)) + equi_enc_feat0 = x + + x = self.equi_encoder.maxpool(x) + equi_enc_feat1 = self.equi_encoder.layer1(x) + equi_enc_feat2 = self.equi_encoder.layer2(equi_enc_feat1) + equi_enc_feat3 = self.equi_encoder.layer3(equi_enc_feat2) + equi_enc_feat4 = self.equi_encoder.layer4(equi_enc_feat3) + + # cube image encoding + cube_inputs = torch.cat( + torch.split(input_cube_image, self.cube_h, dim=-1), dim=0) + + if self.num_layers < 18: + cube_enc_feat0, cube_enc_feat1, cube_enc_feat2, cube_enc_feat3, cube_enc_feat4 \ + = self.cube_encoder(cube_inputs) + else: + + x = self.cube_encoder.conv1(cube_inputs) + x = self.cube_encoder.relu(self.cube_encoder.bn1(x)) + cube_enc_feat0 = x + + x = self.cube_encoder.maxpool(x) + + cube_enc_feat1 = self.cube_encoder.layer1(x) + cube_enc_feat2 = self.cube_encoder.layer2(cube_enc_feat1) + cube_enc_feat3 = self.cube_encoder.layer3(cube_enc_feat2) + cube_enc_feat4 = self.cube_encoder.layer4(cube_enc_feat3) + + # euqi image decoding fused with cubemap features + outputs = {} + + cube_enc_feat4 = torch.cat( + torch.split(cube_enc_feat4, input_equi_image.shape[0], dim=0), + dim=-1) + c2e_enc_feat4 = self.c2e['5'](cube_enc_feat4) + fused_feat4 = self.equi_dec_convs['fusion_5'](equi_enc_feat4, + c2e_enc_feat4) + equi_x = upsample(self.equi_dec_convs['upconv_5'](fused_feat4)) + + cube_enc_feat3 = torch.cat( + torch.split(cube_enc_feat3, input_equi_image.shape[0], dim=0), + dim=-1) + c2e_enc_feat3 = self.c2e['4'](cube_enc_feat3) + fused_feat3 = self.equi_dec_convs['fusion_4'](equi_enc_feat3, + c2e_enc_feat3) + equi_x = torch.cat([equi_x, fused_feat3], 1) + equi_x = self.equi_dec_convs['deconv_4'](equi_x) + equi_x = upsample(self.equi_dec_convs['upconv_4'](equi_x)) + + cube_enc_feat2 = torch.cat( + torch.split(cube_enc_feat2, input_equi_image.shape[0], dim=0), + dim=-1) + c2e_enc_feat2 = self.c2e['3'](cube_enc_feat2) + fused_feat2 = self.equi_dec_convs['fusion_3'](equi_enc_feat2, + c2e_enc_feat2) + equi_x = torch.cat([equi_x, fused_feat2], 1) + equi_x = self.equi_dec_convs['deconv_3'](equi_x) + equi_x = upsample(self.equi_dec_convs['upconv_3'](equi_x)) + + cube_enc_feat1 = torch.cat( + torch.split(cube_enc_feat1, input_equi_image.shape[0], dim=0), + dim=-1) + c2e_enc_feat1 = self.c2e['2'](cube_enc_feat1) + fused_feat1 = self.equi_dec_convs['fusion_2'](equi_enc_feat1, + c2e_enc_feat1) + equi_x = torch.cat([equi_x, fused_feat1], 1) + equi_x = self.equi_dec_convs['deconv_2'](equi_x) + equi_x = upsample(self.equi_dec_convs['upconv_2'](equi_x)) + + cube_enc_feat0 = torch.cat( + torch.split(cube_enc_feat0, input_equi_image.shape[0], dim=0), + dim=-1) + c2e_enc_feat0 = self.c2e['1'](cube_enc_feat0) + fused_feat0 = self.equi_dec_convs['fusion_1'](equi_enc_feat0, + c2e_enc_feat0) + equi_x = torch.cat([equi_x, fused_feat0], 1) + equi_x = self.equi_dec_convs['deconv_1'](equi_x) + equi_x = upsample(self.equi_dec_convs['upconv_1'](equi_x)) + + equi_x = self.equi_dec_convs['deconv_0'](equi_x) + + equi_depth = self.equi_dec_convs['depthconv_0'](equi_x) + outputs['pred_depth'] = self.max_depth * self.sigmoid(equi_depth) + + return outputs diff --git a/modelscope/models/cv/panorama_depth_estimation/networks/util.py b/modelscope/models/cv/panorama_depth_estimation/networks/util.py new file mode 100644 index 00000000..546ccc92 --- /dev/null +++ b/modelscope/models/cv/panorama_depth_estimation/networks/util.py @@ -0,0 +1,111 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import cv2 +import numpy as np +from scipy.ndimage import map_coordinates + + +class Equirec2Cube: + + def __init__(self, equ_h, equ_w, face_w): + ''' + equ_h: int, height of the equirectangular image + equ_w: int, width of the equirectangular image + face_w: int, the length of each face of the cubemap + ''' + + self.equ_h = equ_h + self.equ_w = equ_w + self.face_w = face_w + + self._xyzcube() + self._xyz2coor() + + # For convert R-distance to Z-depth for CubeMaps + cosmap = 1 / np.sqrt((2 * self.grid[..., 0])**2 + + (2 * self.grid[..., 1])**2 + 1) + self.cosmaps = np.concatenate(6 * [cosmap], axis=1)[..., np.newaxis] + + def _xyzcube(self): + ''' + Compute the xyz cordinates of the unit cube in [F R B L U D] format. + ''' + self.xyz = np.zeros((self.face_w, self.face_w * 6, 3), np.float32) + rng = np.linspace(-0.5, 0.5, num=self.face_w, dtype=np.float32) + self.grid = np.stack(np.meshgrid(rng, -rng), -1) + + # Front face (z = 0.5) + self.xyz[:, 0 * self.face_w:1 * self.face_w, [0, 1]] = self.grid + self.xyz[:, 0 * self.face_w:1 * self.face_w, 2] = 0.5 + + # Right face (x = 0.5) + self.xyz[:, 1 * self.face_w:2 * self.face_w, + [2, 1]] = self.grid[:, ::-1] + self.xyz[:, 1 * self.face_w:2 * self.face_w, 0] = 0.5 + + # Back face (z = -0.5) + self.xyz[:, 2 * self.face_w:3 * self.face_w, + [0, 1]] = self.grid[:, ::-1] + self.xyz[:, 2 * self.face_w:3 * self.face_w, 2] = -0.5 + + # Left face (x = -0.5) + self.xyz[:, 3 * self.face_w:4 * self.face_w, [2, 1]] = self.grid + self.xyz[:, 3 * self.face_w:4 * self.face_w, 0] = -0.5 + + # Up face (y = 0.5) + self.xyz[:, 4 * self.face_w:5 * self.face_w, + [0, 2]] = self.grid[::-1, :] + self.xyz[:, 4 * self.face_w:5 * self.face_w, 1] = 0.5 + + # Down face (y = -0.5) + self.xyz[:, 5 * self.face_w:6 * self.face_w, [0, 2]] = self.grid + self.xyz[:, 5 * self.face_w:6 * self.face_w, 1] = -0.5 + + def _xyz2coor(self): + + # x, y, z to longitude and latitude + x, y, z = np.split(self.xyz, 3, axis=-1) + lon = np.arctan2(x, z) + c = np.sqrt(x**2 + z**2) + lat = np.arctan2(y, c) + + # longitude and latitude to equirectangular coordinate + self.coor_x = (lon / (2 * np.pi) + 0.5) * self.equ_w - 0.5 + self.coor_y = (-lat / np.pi + 0.5) * self.equ_h - 0.5 + + def sample_equirec(self, e_img, order=0): + pad_u = np.roll(e_img[[0]], self.equ_w // 2, 1) + pad_d = np.roll(e_img[[-1]], self.equ_w // 2, 1) + e_img = np.concatenate([e_img, pad_d, pad_u], 0) + + return map_coordinates( + e_img, [self.coor_y, self.coor_x], order=order, mode='wrap')[..., + 0] + + def run(self, equ_img, equ_dep=None): + + h, w = equ_img.shape[:2] + if h != self.equ_h or w != self.equ_w: + equ_img = cv2.resize(equ_img, (self.equ_w, self.equ_h)) + if equ_dep is not None: + equ_dep = cv2.resize( + equ_dep, (self.equ_w, self.equ_h), + interpolation=cv2.INTER_NEAREST) + + cube_img = np.stack([ + self.sample_equirec(equ_img[..., i], order=1) + for i in range(equ_img.shape[2]) + ], + axis=-1) # noqa + + if equ_dep is not None: + cube_dep = np.stack([ + self.sample_equirec(equ_dep[..., i], order=0) + for i in range(equ_dep.shape[2]) + ], + axis=-1) # noqa + cube_dep = cube_dep * self.cosmaps + + if equ_dep is not None: + return cube_img, cube_dep + else: + return cube_img diff --git a/modelscope/models/cv/panorama_depth_estimation/unifuse_model.py b/modelscope/models/cv/panorama_depth_estimation/unifuse_model.py new file mode 100644 index 00000000..5fe33e52 --- /dev/null +++ b/modelscope/models/cv/panorama_depth_estimation/unifuse_model.py @@ -0,0 +1,89 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os.path as osp + +import numpy as np +import torch +from torchvision import transforms + +from modelscope.metainfo import Models +from modelscope.models.base.base_torch_model import TorchModel +from modelscope.models.builder import MODELS +from modelscope.models.cv.panorama_depth_estimation.networks import (Equi, + UniFuse) +from modelscope.models.cv.panorama_depth_estimation.networks.util import \ + Equirec2Cube +from modelscope.outputs import OutputKeys +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.logger import get_logger + +logger = get_logger() + + +@MODELS.register_module( + Tasks.panorama_depth_estimation, + module_name=Models.unifuse_depth_estimation) +class PanoramaDepthEstimation(TorchModel): + """ + UniFuse: Unidirectional Fusion for 360 Panorama Depth Estimation + https://arxiv.org/abs/2102.03550 + """ + + def __init__(self, model_dir: str, **kwargs): + """ + Args: + model_dir: the path of the pretrained model file + """ + super().__init__(model_dir, **kwargs) + self.device = torch.device( + 'cuda' if torch.cuda.is_available() else 'cpu') + + # load model + model_path = osp.join(model_dir, ModelFile.TORCH_MODEL_FILE) + logger.info(f'loading model {model_path}') + model_dict = torch.load(model_path, map_location=torch.device('cpu')) + Net_dict = {'UniFuse': UniFuse, 'Equi': Equi} + Net = Net_dict[model_dict['net']] + self.w = model_dict['width'] + self.h = model_dict['height'] + self.max_depth_meters = 10.0 + self.e2c = Equirec2Cube(self.h, self.w, self.h // 2) + self.to_tensor = transforms.ToTensor() + self.normalize = transforms.Normalize( + mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + + # build model + self.model = Net( + model_dict['layers'], + model_dict['height'], + model_dict['width'], + max_depth=self.max_depth_meters, + fusion_type=model_dict['fusion'], + se_in_fusion=model_dict['se_in_fusion']) + + # load state dict + self.model.to(self.device) + model_state_dict = self.model.state_dict() + self.model.load_state_dict( + {k: v + for k, v in model_dict.items() if k in model_state_dict}) + self.model.eval() + + logger.info(f'model init done! Device:{self.device}') + + def forward(self, Inputs): + """ + Args: + Inputs: model inputs containning equirectangular panorama images and the corresponding cubmap images + The torch size of Inputs['rgb'] should be [n, 3, 512, 1024] + The torch size of Inputs['cube_rgb'] should be [n, 3, 256, 1536] + Returns: + Unifuse model outputs containing the predicted equirectangular depth images in metric + """ + equi_inputs = Inputs['rgb'].to(self.device) + cube_inputs = Inputs['cube_rgb'].to(self.device) + return self.model(equi_inputs, cube_inputs) + + def postprocess(self, Inputs): + depth_result = Inputs['pred_depth'][0] + results = {OutputKeys.DEPTHS: depth_result} + return results diff --git a/modelscope/pipelines/builder.py b/modelscope/pipelines/builder.py index 05859683..2514b367 100644 --- a/modelscope/pipelines/builder.py +++ b/modelscope/pipelines/builder.py @@ -176,6 +176,9 @@ DEFAULT_MODEL_FOR_PIPELINE = { Tasks.video_depth_estimation: (Pipelines.video_depth_estimation, 'damo/cv_dro-resnet18_video-depth-estimation_indoor'), + Tasks.panorama_depth_estimation: + (Pipelines.panorama_depth_estimation, + 'damo/cv_unifuse_panorama-depth-estimation'), Tasks.image_style_transfer: (Pipelines.image_style_transfer, 'damo/cv_aams_style-transfer_damo'), Tasks.face_image_generation: (Pipelines.face_image_generation, @@ -247,9 +250,9 @@ DEFAULT_MODEL_FOR_PIPELINE = { 'damo/cv_googlenet_pgl-video-summarization'), Tasks.image_skychange: (Pipelines.image_skychange, 'damo/cv_hrnetocr_skychange'), - Tasks.translation_evaluation: - (Pipelines.translation_evaluation, - 'damo/nlp_unite_mup_translation_evaluation_multilingual_large'), + Tasks.translation_evaluation: ( + Pipelines.translation_evaluation, + 'damo/nlp_unite_mup_translation_evaluation_multilingual_large'), Tasks.video_object_segmentation: ( Pipelines.video_object_segmentation, 'damo/cv_rdevos_video-object-segmentation'), diff --git a/modelscope/pipelines/cv/__init__.py b/modelscope/pipelines/cv/__init__.py index 3358c961..3e13b528 100644 --- a/modelscope/pipelines/cv/__init__.py +++ b/modelscope/pipelines/cv/__init__.py @@ -76,6 +76,7 @@ if TYPE_CHECKING: from .pointcloud_sceneflow_estimation_pipeline import PointCloudSceneFlowEstimationPipeline from .maskdino_instance_segmentation_pipeline import MaskDINOInstanceSegmentationPipeline from .image_mvs_depth_estimation_pipeline import ImageMultiViewDepthEstimationPipeline + from .panorama_depth_estimation_pipeline import PanoramaDepthEstimationPipeline from .ddcolor_image_colorization_pipeline import DDColorImageColorizationPipeline else: diff --git a/modelscope/pipelines/cv/panorama_depth_estimation_pipeline.py b/modelscope/pipelines/cv/panorama_depth_estimation_pipeline.py new file mode 100644 index 00000000..a1973285 --- /dev/null +++ b/modelscope/pipelines/cv/panorama_depth_estimation_pipeline.py @@ -0,0 +1,85 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import Any, Dict, Union + +import cv2 +import numpy as np +import PIL +import torch + +from modelscope.metainfo import Pipelines +from modelscope.outputs import OutputKeys +from modelscope.pipelines.base import Input, Model, Pipeline +from modelscope.pipelines.builder import PIPELINES +from modelscope.preprocessors import LoadImage +from modelscope.utils.constant import Tasks +from modelscope.utils.cv.image_utils import depth_to_color +from modelscope.utils.logger import get_logger + +logger = get_logger() + + +@PIPELINES.register_module( + Tasks.panorama_depth_estimation, + module_name=Pipelines.panorama_depth_estimation) +class PanoramaDepthEstimationPipeline(Pipeline): + """ This pipeline will estimation the depth panoramic image from one rgb panoramic image. + The input panoramic image should be equirectanlar, in the size of 512x1024. + + Example: + '''python + >>> import cv2 + >>> from modelscope.outputs import OutputKeys + >>> from modelscope.pipelines import pipeline + >>> from modelscope.utils.constant import Tasks + + >>> task = 'panorama-depth-estimation' + >>> model_id = 'damo/cv_unifuse_image-depth-estimation' + + >>> input_location = 'data/test/images/panorama_depth_estimation.jpg' + >>> estimator = pipeline(Tasks.panorama_depth_estimation, model=model_id) + >>> result = estimator(input_location) + >>> depth_vis = result[OutputKeys.DEPTHS_COLOR] + >>> cv2.imwrite('result.jpg', depth_vis) + ''' + """ + + def __init__(self, model: str, **kwargs): + """ + use `model` to create a panorama depth estimation pipeline for prediction + Args: + model: model id on modelscope hub. + """ + super().__init__(model=model, **kwargs) + + logger.info('depth estimation model, pipeline init') + + def preprocess(self, input: Input) -> Dict[str, Any]: + img = LoadImage.convert_to_ndarray(input) + H, W = 512, 1024 + img = cv2.resize(img, dsize=(W, H), interpolation=cv2.INTER_CUBIC) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + cube_img = self.model.e2c.run(img) + data = {} + rgb = self.model.to_tensor(img.copy()) + cube_rgb = self.model.to_tensor(cube_img.copy()) + rgb = self.model.normalize(rgb) + cube_rgb = self.model.normalize(cube_rgb) + data['rgb'] = rgb[None, ...] + data['cube_rgb'] = cube_rgb[None, ...] + return data + + def forward(self, input: Dict[str, Any]) -> Dict[str, Any]: + results = self.model.forward(input) + return results + + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + results = self.model.postprocess(inputs) + depths = results[OutputKeys.DEPTHS] + if isinstance(depths, torch.Tensor): + depths = depths.detach().cpu().squeeze().numpy() + depths_color = depth_to_color(depths) + outputs = { + OutputKeys.DEPTHS: depths, + OutputKeys.DEPTHS_COLOR: depths_color + } + return outputs diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index f73ec428..96b6d1b1 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -51,6 +51,7 @@ class CVTasks(object): semantic_segmentation = 'semantic-segmentation' image_depth_estimation = 'image-depth-estimation' video_depth_estimation = 'video-depth-estimation' + panorama_depth_estimation = 'panorama-depth-estimation' portrait_matting = 'portrait-matting' text_driven_segmentation = 'text-driven-segmentation' shop_segmentation = 'shop-segmentation' diff --git a/tests/pipelines/test_panorama_depth_estimation.py b/tests/pipelines/test_panorama_depth_estimation.py new file mode 100644 index 00000000..99e575e3 --- /dev/null +++ b/tests/pipelines/test_panorama_depth_estimation.py @@ -0,0 +1,34 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import unittest + +import cv2 +import numpy as np + +from modelscope.outputs import OutputKeys +from modelscope.pipelines import pipeline +from modelscope.utils.constant import Tasks +from modelscope.utils.cv.image_utils import depth_to_color +from modelscope.utils.demo_utils import DemoCompatibilityCheck +from modelscope.utils.test_utils import test_level + + +class PanoramaDepthEstimationTest(unittest.TestCase, DemoCompatibilityCheck): + + def setUp(self) -> None: + self.task = 'panorama-depth-estimation' + self.model_id = 'damo/cv_unifuse_panorama-depth-estimation' + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_panorama_depth_estimation(self): + input_location = 'data/test/images/panorama_depth_estimation.jpg' + estimator = pipeline( + Tasks.panorama_depth_estimation, model=self.model_id) + result = estimator(input_location) + depth_vis = result[OutputKeys.DEPTHS_COLOR] + cv2.imwrite('result.jpg', depth_vis) + print('test_panorama_depth_estimation DONE') + + +if __name__ == '__main__': + unittest.main()