add panorama_depth_estimation

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

data/test/images/panorama_depth_estimation.jpg

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0df5f2de59df6b55d8ee5d414cc2f98d714e14b518c159d4085ad2ac65d36627
+size 137606

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'

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

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={},
+    )

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

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

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

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" <https://arxiv.org/abs/1801.04381>`_.
+
+    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

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" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    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" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    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" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    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" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    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" <https://arxiv.org/pdf/1512.03385.pdf>`_
+
+    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" <https://arxiv.org/pdf/1611.05431.pdf>`_
+
+    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" <https://arxiv.org/pdf/1611.05431.pdf>`_
+
+    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" <https://arxiv.org/pdf/1605.07146.pdf>`_
+
+    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" <https://arxiv.org/pdf/1605.07146.pdf>`_
+
+    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)

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

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

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

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'),

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:

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

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'

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()