An Out-Of-The-Box Version of Metric3D
if __name__=='__main__':
    import os
    import os.path as osp
    import numpy as np
    import cv2
    import json

    code_root = '/mnt/nas/share/home/xugk/MetricDepth_test/'

    data_root = osp.join(code_root, 'data/kitti_demo')
    split_root = code_root

    files = []
    rgb_root = osp.join(data_root, 'rgb')
    depth_root = osp.join(data_root, 'depth')
    for rgb_file in os.listdir(rgb_root):
        rgb_path = osp.join(rgb_root, rgb_file).split(split_root)[-1]
        depth_path = rgb_path.replace('/rgb/', '/depth/')
        cam_in = [707.0493, 707.0493, 604.0814, 180.5066]
        depth_scale = 256.

        meta_data = {}
        meta_data['cam_in'] = cam_in
        meta_data['rgb'] = rgb_path
        meta_data['depth'] = depth_path
        meta_data['depth_scale'] = depth_scale
        files.append(meta_data)
    files_dict = dict(files=files)

    with open(osp.join(code_root, 'data/kitti_demo/test_annotations.json'), 'w') as f:
        json.dump(files_dict, f)
        
data/gene_annos_kitti_demo.py
hugoycj/Metric3D_MVP
if __name__=='__main__':
    import os
    import os.path as osp
    import numpy as np
    import cv2
    import json

    code_root = '/mnt/nas/share/home/xugk/MetricDepth_test/'

    data_root = osp.join(code_root, 'data/nyu_demo')
    split_root = code_root

    files = []
    rgb_root = osp.join(data_root, 'rgb')
    depth_root = osp.join(data_root, 'depth')
    for rgb_file in os.listdir(rgb_root):
        rgb_path = osp.join(rgb_root, rgb_file).split(split_root)[-1]
        depth_path = rgb_path.replace('.jpg', '.png').replace('/rgb_', '/sync_depth_').replace('/rgb/', '/depth/')
        cam_in = [518.8579, 519.46961, 325.58245, 253.73617]
        depth_scale = 1000.

        meta_data = {}
        meta_data['cam_in'] = cam_in
        meta_data['rgb'] = rgb_path
        meta_data['depth'] = depth_path
        meta_data['depth_scale'] = depth_scale
        files.append(meta_data)
    files_dict = dict(files=files)

    with open(osp.join(code_root, 'data/nyu_demo/test_annotations.json'), 'w') as f:
        json.dump(files_dict, f)
data/gene_annos_nyu_demo.py
hugoycj/Metric3D_MVP
from .public_datasets import *
from .pretrained_weight import *
data_info/__init__.py
hugoycj/Metric3D_MVP

mldb_info={}

mldb_info['checkpoint']={
    'mldb_root': '/mnt/nas/share/home/xugk/ckpt', # NOTE: modify it to the pretrained ckpt root

    # pretrained weight for convnext
    'convnext_tiny': 'convnext/convnext_tiny_22k_1k_384.pth',
    'convnext_small': 'convnext/convnext_small_22k_1k_384.pth',
    'convnext_base': 'convnext/convnext_base_22k_1k_384.pth',
    'convnext_large': 'convnext/convnext_large_22k_1k_384.pth',
}
data_info/pretrained_weight.py
hugoycj/Metric3D_MVP
mldb_info = {}

mldb_info['NYU']={
    'mldb_root': '/mnt/nas/share/home/xugk/data/',
    'data_root': 'nyu',
    'test_annotations_path': 'nyu/test_annotation.json',
}

data_info/public_datasets.py
hugoycj/Metric3D_MVP
dependencies = ['torch']
from mono.utils.logger import setup_logger
from mono.utils.comm import init_env
from mono.model.monodepth_model import get_configured_monodepth_model
from mono.utils.running import load_ckpt
from mono.utils.do_test import do_scalecano_test_with_custom_data
from mono.utils.mldb import load_data_info, reset_ckpt_path
from mono.utils.custom_data import load_from_annos, load_data
import torch

def Metric3D(pretrained=False, **kwargs):
    """
    # This docstring shows up in hub.help()
    MonoDepth model
    pretrained (bool): kwargs, load pretrained weights into the model
    """
    # monodepth model
    model = get_configured_monodepth_model(**kwargs)
    if pretrained:
        ckpt = torch.hub.load_state_dict_from_url(
            'https://path/to/pretrained/weights',
            progress=True,
            check_hash=True
        )
        model.load_state_dict(ckpt)
    return model
hubconf.py
hugoycj/Metric3D_MVP
_base_=[
       '../_base_/models/encoder_decoder/convnext_large.hourglassdecoder.py',
       '../_base_/datasets/_data_base_.py',
       '../_base_/default_runtime.py',
       ]

model = dict(
    backbone=dict(
        pretrained=False,
    )
)

# configs of the canonical space
data_basic=dict(
    canonical_space = dict(
        img_size=(512, 960),
        focal_length=1000.0,
    ),
    depth_range=(0, 1),
    depth_normalize=(0.3, 150),
    crop_size = (544, 1216),
) 

batchsize_per_gpu = 2
thread_per_gpu = 4

mono/configs/HourglassDecoder/convlarge.0.3_150.py
hugoycj/Metric3D_MVP
_base_=[
       '../_base_/models/encoder_decoder/convnext_large.hourglassdecoder.py',
       '../_base_/datasets/_data_base_.py',
       '../_base_/default_runtime.py',
       ]

model = dict(
    backbone=dict(
        pretrained=False,
    )
)

# configs of the canonical space
data_basic=dict(
    canonical_space = dict(
        img_size=(512, 960),
        focal_length=1000.0,
    ),
    depth_range=(0, 1),
    depth_normalize=(0.3, 150),
    crop_size = (512, 1088),
) 

batchsize_per_gpu = 2
thread_per_gpu = 4

mono/configs/HourglassDecoder/test_kitti_convlarge.0.3_150.py
hugoycj/Metric3D_MVP
_base_=[
       '../_base_/models/encoder_decoder/convnext_large.hourglassdecoder.py',
       '../_base_/datasets/_data_base_.py',
       '../_base_/default_runtime.py',
       ]

model = dict(
    backbone=dict(
        pretrained=False,
    )
)

# configs of the canonical space
data_basic=dict(
    canonical_space = dict(
        img_size=(512, 960),
        focal_length=1000.0,
    ),
    depth_range=(0, 1),
    depth_normalize=(0.3, 150),
    crop_size = (480, 1216),
) 

batchsize_per_gpu = 2
thread_per_gpu = 4

mono/configs/HourglassDecoder/test_nyu_convlarge.0.3_150.py
hugoycj/Metric3D_MVP
# canonical camera setting and basic data setting
#  
data_basic=dict(
    canonical_space = dict(
        img_size=(540, 960),
        focal_length=1196.0,
    ),
    depth_range=(0.9, 150),
    depth_normalize=(0.006, 1.001),
    crop_size = (512, 960),
    clip_depth_range=(0.9, 150),
)    

mono/configs/_base_/datasets/_data_base_.py
hugoycj/Metric3D_MVP

load_from = None
cudnn_benchmark = True
test_metrics = ['abs_rel', 'rmse', 'silog', 'delta1', 'delta2', 'delta3','rmse_log', 'log10', 'sq_rel']

mono/configs/_base_/default_runtime.py
hugoycj/Metric3D_MVP
#_base_ = ['./_model_base_.py',]

#'https://download.openmmlab.com/mmclassification/v0/convnext/downstream/convnext-large_3rdparty_in21k_20220301-e6e0ea0a.pth' 
model = dict(
    #type='EncoderDecoderAuxi',
    backbone=dict(
        type='convnext_large',
        pretrained=True, 
        in_22k=True,
        out_indices=[0, 1, 2, 3],
        drop_path_rate=0.4,
        layer_scale_init_value=1.0,
        checkpoint='data/pretrained_weight_repo/convnext/convnext_large_22k_1k_384.pth',
        prefix='backbones.',
        out_channels=[192, 384, 768, 1536]),
    )

mono/configs/_base_/models/backbones/convnext_large.py
hugoycj/Metric3D_MVP
# model settings
_base_ = ['../backbones/convnext_large.py',]
model = dict(
    type='DensePredModel',
    decode_head=dict(
        type='HourglassDecoder',
        in_channels=[192, 384, 768, 1536],
        decoder_channel=[128, 128, 256, 512],
        prefix='decode_heads.'),
)

mono/configs/_base_/models/encoder_decoder/convnext_large.hourglassdecoder.py
hugoycj/Metric3D_MVP
from .monodepth_model import DepthModel
# from .__base_model__ import BaseDepthModel


__all__ = ['DepthModel', 'BaseDepthModel']

mono/model/__init__.py
hugoycj/Metric3D_MVP
import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import trunc_normal_, DropPath
from timm.models.registry import register_model

class Block(nn.Module):
    r""" ConvNeXt Block. There are two equivalent implementations:
    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
    We use (2) as we find it slightly faster in PyTorch
    
    Args:
        dim (int): Number of input channels.
        drop_path (float): Stochastic depth rate. Default: 0.0
        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
    """
    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6):
        super().__init__()
        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
        self.norm = LayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(4 * dim, dim)
        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), 
                                    requires_grad=True) if layer_scale_init_value > 0 else None
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        input = x
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)

        x = input + self.drop_path(x)
        return x

class ConvNeXt(nn.Module):
    r""" ConvNeXt
        A PyTorch impl of : `A ConvNet for the 2020s`  -
          https://arxiv.org/pdf/2201.03545.pdf
    Args:
        in_chans (int): Number of input image channels. Default: 3
        num_classes (int): Number of classes for classification head. Default: 1000
        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
        drop_path_rate (float): Stochastic depth rate. Default: 0.
        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
    """
    def __init__(self, in_chans=3, num_classes=1000, 
                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], drop_path_rate=0., 
                 layer_scale_init_value=1e-6, head_init_scale=1.,
                 **kwargs,):
        super().__init__()

        self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers
        stem = nn.Sequential(
            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
            LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
        )
        self.downsample_layers.append(stem)
        for i in range(3):
            downsample_layer = nn.Sequential(
                    LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
                    nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2),
            )
            self.downsample_layers.append(downsample_layer)

        self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks
        dp_rates=[x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] 
        cur = 0
        for i in range(4):
            stage = nn.Sequential(
                *[Block(dim=dims[i], drop_path=dp_rates[cur + j], 
                layer_scale_init_value=layer_scale_init_value) for j in range(depths[i])]
            )
            self.stages.append(stage)
            cur += depths[i]

        #self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer
        #self.head = nn.Linear(dims[-1], num_classes)

        self.apply(self._init_weights)
        #self.head.weight.data.mul_(head_init_scale)
        #self.head.bias.data.mul_(head_init_scale)

    def _init_weights(self, m):
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            trunc_normal_(m.weight, std=.02)
            nn.init.constant_(m.bias, 0)

    def forward_features(self, x):
        features = []
        for i in range(4):
            x = self.downsample_layers[i](x)
            x = self.stages[i](x)
            features.append(x)
        return features # global average pooling, (N, C, H, W) -> (N, C)

    def forward(self, x):
        #x = self.forward_features(x)
        #x = self.head(x)
        features = self.forward_features(x)
        return features

class LayerNorm(nn.Module):
    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 
    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
    shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
    with shape (batch_size, channels, height, width).
    """
    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ["channels_last", "channels_first"]:
            raise NotImplementedError 
        self.normalized_shape = (normalized_shape, )
    
    def forward(self, x):
        if self.data_format == "channels_last":
            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
        elif self.data_format == "channels_first":
            u = x.mean(1, keepdim=True)
            s = (x - u).pow(2).mean(1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x


model_urls = {
    "convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
    "convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
    "convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
    "convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
    "convnext_tiny_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_224.pth",
    "convnext_small_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth",
    "convnext_base_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth",
    "convnext_large_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth",
    "convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth",
}

def convnext_tiny(pretrained=True,in_22k=False, **kwargs):
    model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
    if pretrained:
        checkpoint = torch.load(kwargs['checkpoint'], map_location="cpu")
        #url = model_urls['convnext_tiny_22k'] if in_22k else model_urls['convnext_tiny_1k']
        #checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
        model_dict = model.state_dict()
        pretrained_dict = {}
        unmatched_pretrained_dict = {}
        for k, v in checkpoint['model'].items():
            if k in model_dict:
                pretrained_dict[k] = v
            else:
                unmatched_pretrained_dict[k] = v
        model_dict.update(pretrained_dict)
        model.load_state_dict(model_dict)
        print(
            'Successfully loaded pretrained %d params, and %d paras are unmatched.'
            %(len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys())))
        print('Unmatched pretrained paras are :', unmatched_pretrained_dict.keys())
    return model

def convnext_small(pretrained=True,in_22k=False, **kwargs):
    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
    if pretrained:
        checkpoint = torch.load(kwargs['checkpoint'], map_location="cpu")
        #url = model_urls['convnext_small_22k'] if in_22k else model_urls['convnext_small_1k']
        #checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
        model_dict = model.state_dict()
        pretrained_dict = {}
        unmatched_pretrained_dict = {}
        for k, v in checkpoint['model'].items():
            if k in model_dict:
                pretrained_dict[k] = v
            else:
                unmatched_pretrained_dict[k] = v
        model_dict.update(pretrained_dict)
        model.load_state_dict(model_dict)
        print(
            'Successfully loaded pretrained %d params, and %d paras are unmatched.'
            %(len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys())))
        print('Unmatched pretrained paras are :', unmatched_pretrained_dict.keys())
    return model

def convnext_base(pretrained=True, in_22k=False, **kwargs):
    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
    if pretrained:
        checkpoint = torch.load(kwargs['checkpoint'], map_location="cpu")
        #url = model_urls['convnext_base_22k'] if in_22k else model_urls['convnext_base_1k']
        #checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
        model_dict = model.state_dict()
        pretrained_dict = {}
        unmatched_pretrained_dict = {}
        for k, v in checkpoint['model'].items():
            if k in model_dict:
                pretrained_dict[k] = v
            else:
                unmatched_pretrained_dict[k] = v
        model_dict.update(pretrained_dict)
        model.load_state_dict(model_dict)
        print(
            'Successfully loaded pretrained %d params, and %d paras are unmatched.'
            %(len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys())))
        print('Unmatched pretrained paras are :', unmatched_pretrained_dict.keys())
    return model

def convnext_large(pretrained=True, in_22k=False, **kwargs):
    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
    if pretrained:
        checkpoint = torch.load(kwargs['checkpoint'], map_location="cpu")
        #url = model_urls['convnext_large_22k'] if in_22k else model_urls['convnext_large_1k']
        #checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
        model_dict = model.state_dict()
        pretrained_dict = {}
        unmatched_pretrained_dict = {}
        for k, v in checkpoint['model'].items():
            if k in model_dict:
                pretrained_dict[k] = v
            else:
                unmatched_pretrained_dict[k] = v
        model_dict.update(pretrained_dict)
        model.load_state_dict(model_dict)
        print(
            'Successfully loaded pretrained %d params, and %d paras are unmatched.'
            %(len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys())))
        print('Unmatched pretrained paras are :', unmatched_pretrained_dict.keys())
    return model

def convnext_xlarge(pretrained=True, in_22k=False, **kwargs):
    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
    if pretrained:
        assert in_22k, "only ImageNet-22K pre-trained ConvNeXt-XL is available; please set in_22k=True"
        checkpoint = torch.load(kwargs['checkpoint'], map_location="cpu")
        #url = model_urls['convnext_xlarge_22k']
        #checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
        model_dict = model.state_dict()
        pretrained_dict = {}
        unmatched_pretrained_dict = {}
        for k, v in checkpoint['model'].items():
            if k in model_dict:
                pretrained_dict[k] = v
            else:
                unmatched_pretrained_dict[k] = v
        model_dict.update(pretrained_dict)
        model.load_state_dict(model_dict)
        print(
            'Successfully loaded pretrained %d params, and %d paras are unmatched.'
            %(len(pretrained_dict.keys()), len(unmatched_pretrained_dict.keys())))
        print('Unmatched pretrained paras are :', unmatched_pretrained_dict.keys())
    return model

if __name__ == '__main__':
    import torch
    model = convnext_base(True, in_22k=False).cuda()

    rgb = torch.rand((2, 3, 256, 256)).cuda()
    out = model(rgb)
    print(len(out))
    for i, ft in enumerate(out):
        print(i, ft.shape)

mono/model/backbones/ConvNeXt.py
hugoycj/Metric3D_MVP
from .ConvNeXt import convnext_xlarge
from .ConvNeXt import convnext_small
from .ConvNeXt import convnext_base
from .ConvNeXt import convnext_large
from .ConvNeXt import convnext_tiny


__all__ = [
    'convnext_xlarge', 'convnext_small', 'convnext_base', 'convnext_large', 'convnext_tiny', 
]

mono/model/backbones/__init__.py
hugoycj/Metric3D_MVP
import torch
import torch.nn as nn
import numpy as np
import math
import torch.nn.functional as F

def compute_depth_expectation(prob, depth_values):
    depth_values = depth_values.view(*depth_values.shape, 1, 1)
    depth = torch.sum(prob * depth_values, 1)
    return depth

class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        super(ConvBlock, self).__init__()

        if kernel_size == 3:
            self.conv = nn.Sequential(
                nn.ReflectionPad2d(1),
                nn.Conv2d(in_channels, out_channels, 3, padding=0, stride=1),
            )
        elif kernel_size == 1:
            self.conv = nn.Conv2d(int(in_channels), int(out_channels), 1, padding=0, stride=1)

        self.nonlin = nn.ELU(inplace=True)

    def forward(self, x):
        out = self.conv(x)
        out = self.nonlin(out)
        return out


class ConvBlock_double(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        super(ConvBlock_double, self).__init__()

        if kernel_size == 3:
            self.conv = nn.Sequential(
                nn.ReflectionPad2d(1),
                nn.Conv2d(in_channels, out_channels, 3, padding=0, stride=1),
            )
        elif kernel_size == 1:
            self.conv = nn.Conv2d(int(in_channels), int(out_channels), 1, padding=0, stride=1)

        self.nonlin = nn.ELU(inplace=True)
        self.conv_2 = nn.Conv2d(out_channels, out_channels, 1, padding=0, stride=1)
        self.nonlin_2 =nn.ELU(inplace=True)

    def forward(self, x):
        out = self.conv(x)
        out = self.nonlin(out)
        out = self.conv_2(out)
        out = self.nonlin_2(out)
        return out
    
class DecoderFeature(nn.Module):
    def __init__(self, feat_channels, num_ch_dec=[64, 64, 128, 256]):
        super(DecoderFeature, self).__init__()
        self.num_ch_dec = num_ch_dec
        self.feat_channels = feat_channels

        self.upconv_3_0 = ConvBlock(self.feat_channels[3], self.num_ch_dec[3], kernel_size=1)
        self.upconv_3_1 = ConvBlock_double(
            self.feat_channels[2] + self.num_ch_dec[3],
            self.num_ch_dec[3],
            kernel_size=1)
        
        self.upconv_2_0 = ConvBlock(self.num_ch_dec[3], self.num_ch_dec[2], kernel_size=3)
        self.upconv_2_1 = ConvBlock_double(
            self.feat_channels[1] + self.num_ch_dec[2],
            self.num_ch_dec[2],
            kernel_size=3)
        
        self.upconv_1_0 = ConvBlock(self.num_ch_dec[2], self.num_ch_dec[1], kernel_size=3)
        self.upconv_1_1 = ConvBlock_double(
            self.feat_channels[0] + self.num_ch_dec[1],
            self.num_ch_dec[1],
            kernel_size=3)
        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')

    def forward(self, ref_feature):
        x = ref_feature[3]

        x = self.upconv_3_0(x)
        x = torch.cat((self.upsample(x), ref_feature[2]), 1)
        x = self.upconv_3_1(x)

        x = self.upconv_2_0(x)
        x = torch.cat((self.upsample(x), ref_feature[1]), 1)
        x = self.upconv_2_1(x)

        x = self.upconv_1_0(x)
        x = torch.cat((self.upsample(x), ref_feature[0]), 1)
        x = self.upconv_1_1(x)
        return x


class UNet(nn.Module):
    def __init__(self, inp_ch=32, output_chal=1, down_sample_times=3, channel_mode='v0'):
        super(UNet, self).__init__()
        basic_block = ConvBnReLU
        num_depth = 128

        self.conv0 = basic_block(inp_ch, num_depth)
        if channel_mode == 'v0':
            channels = [num_depth, num_depth//2, num_depth//4, num_depth//8, num_depth // 8]
        elif channel_mode == 'v1':
            channels = [num_depth, num_depth, num_depth, num_depth, num_depth, num_depth]
        self.down_sample_times = down_sample_times
        for i in range(down_sample_times):
            setattr(
                self, 'conv_%d' % i,
                nn.Sequential(
                    basic_block(channels[i], channels[i+1], stride=2),
                    basic_block(channels[i+1], channels[i+1])
                )
            )
        for i in range(down_sample_times-1,-1,-1):
            setattr(self, 'deconv_%d' % i,
                    nn.Sequential(
                        nn.ConvTranspose2d(
                            channels[i+1],
                            channels[i],
                            kernel_size=3,
                            padding=1,
                            output_padding=1,
                            stride=2,
                            bias=False),
                        nn.BatchNorm2d(channels[i]),
                        nn.ReLU(inplace=True)
                    )
                )
            self.prob = nn.Conv2d(num_depth, output_chal, 1, stride=1, padding=0)
    
    def forward(self, x):
        features = {}
        conv0 = self.conv0(x)
        x = conv0
        features[0] = conv0
        for i in range(self.down_sample_times):
            x = getattr(self, 'conv_%d' % i)(x)
            features[i+1] = x
        for i in range(self.down_sample_times-1,-1,-1):
            x = features[i] + getattr(self, 'deconv_%d' % i)(x)
        x = self.prob(x)
        return x

class ConvBnReLU(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1):
        super(ConvBnReLU, self).__init__()
        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=pad,
            bias=False
        )
        self.bn = nn.BatchNorm2d(out_channels)
    
    def forward(self, x):
        return F.relu(self.bn(self.conv(x)), inplace=True)


class HourglassDecoder(nn.Module):
    def __init__(self, cfg):
        super(HourglassDecoder, self).__init__()
        self.inchannels = cfg.model.decode_head.in_channels #  [256, 512, 1024, 2048]
        self.decoder_channels = cfg.model.decode_head.decoder_channel # [64, 64, 128, 256]
        self.min_val = cfg.data_basic.depth_normalize[0]
        self.max_val = cfg.data_basic.depth_normalize[1]

        self.num_ch_dec = self.decoder_channels # [64, 64, 128, 256]
        self.num_depth_regressor_anchor = 512
        self.feat_channels = self.inchannels
        unet_in_channel = self.num_ch_dec[1]
        unet_out_channel = 256

        self.decoder_mono = DecoderFeature(self.feat_channels, self.num_ch_dec)
        self.conv_out_2 = UNet(inp_ch=unet_in_channel,
                               output_chal=unet_out_channel + 1,
                               down_sample_times=3,
                               channel_mode='v0',
                               )

        self.depth_regressor_2 = nn.Sequential(
            nn.Conv2d(unet_out_channel,
                      self.num_depth_regressor_anchor,
                      kernel_size=3,
                      padding=1,
                ),
            nn.BatchNorm2d(self.num_depth_regressor_anchor),
            nn.ReLU(inplace=True),
            nn.Conv2d(
                self.num_depth_regressor_anchor,
                self.num_depth_regressor_anchor,
                kernel_size=1,
            )
        )
        self.residual_channel = 16
        self.conv_up_2 = nn.Sequential(
            nn.Conv2d(1 + 2 + unet_out_channel, self.residual_channel, 3, padding=1),
            nn.BatchNorm2d(self.residual_channel),
            nn.ReLU(),
            nn.Conv2d(self.residual_channel, self.residual_channel, 3, padding=1),
            nn.Upsample(scale_factor=4),
            nn.Conv2d(self.residual_channel, self.residual_channel, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(self.residual_channel, 1, 1, padding=0),
        )
    
    def get_bins(self, bins_num):
        depth_bins_vec = torch.linspace(math.log(self.min_val), math.log(self.max_val), bins_num, device='cuda')
        depth_bins_vec = torch.exp(depth_bins_vec)
        return depth_bins_vec
    
    def register_depth_expectation_anchor(self, bins_num, B):
        depth_bins_vec = self.get_bins(bins_num)
        depth_bins_vec = depth_bins_vec.unsqueeze(0).repeat(B, 1)
        self.register_buffer('depth_expectation_anchor', depth_bins_vec, persistent=False)

    def upsample(self, x, scale_factor=2):
        return F.interpolate(x, scale_factor=scale_factor, mode='nearest')

    def regress_depth_2(self, feature_map_d):
        prob = self.depth_regressor_2(feature_map_d).softmax(dim=1)
        B = prob.shape[0]
        if "depth_expectation_anchor" not in self._buffers:
            self.register_depth_expectation_anchor(self.num_depth_regressor_anchor, B)
        d = compute_depth_expectation(
            prob,
            self.depth_expectation_anchor[:B, ...]
        ).unsqueeze(1)
        return d

    def create_mesh_grid(self, height, width, batch, device="cuda", set_buffer=True):
        y, x = torch.meshgrid([torch.arange(0, height, dtype=torch.float32, device=device),
                               torch.arange(0, width, dtype=torch.float32, device=device)], indexing='ij')
        meshgrid = torch.stack((x, y))
        meshgrid = meshgrid.unsqueeze(0).repeat(batch, 1, 1, 1)
        return meshgrid

    def forward(self, features_mono, **kwargs):
        '''
        trans_ref2src: list of transformation matrix from the reference view to source view. [B, 4, 4]
        inv_intrinsic_pool: list of inverse intrinsic matrix.
        features_mono: features of reference and source views. [[ref_f1, ref_f2, ref_f3, ref_f4],[src1_f1, src1_f2, src1_f3, src1_f4], ...].
        '''
        outputs = {}
        # get encoder feature of the reference view
        ref_feat = features_mono

        feature_map_mono = self.decoder_mono(ref_feat)
        feature_map_mono_pred = self.conv_out_2(feature_map_mono)
        confidence_map_2 = feature_map_mono_pred[:, -1:, :, :]
        feature_map_d_2 = feature_map_mono_pred[:, :-1, :, :]

        depth_pred_2 = self.regress_depth_2(feature_map_d_2)

        B, _, H, W = depth_pred_2.shape

        meshgrid = self.create_mesh_grid(H, W, B)

        depth_pred_mono = self.upsample(depth_pred_2, scale_factor=4) + 1e-1 * \
            self.conv_up_2(
                torch.cat((depth_pred_2, meshgrid[:B, ...], feature_map_d_2), 1)
            )
        confidence_map_mono = self.upsample(confidence_map_2, scale_factor=4)

        outputs=dict(
            prediction=depth_pred_mono,
            confidence=confidence_map_mono,
            pred_logit=None,
        )
        return outputs
mono/model/decode_heads/HourGlassDecoder.py
hugoycj/Metric3D_MVP
from .HourGlassDecoder import HourglassDecoder

__all__=['HourglassDecoder']

mono/model/decode_heads/__init__.py
hugoycj/Metric3D_MVP
import torch
import torch.nn as nn
from mono.utils.comm import get_func


class BaseDepthModel(nn.Module):
    def __init__(self, cfg, **kwargs) -> None:
        super(BaseDepthModel, self).__init__()
        model_type = cfg.model.type
        self.depth_model = get_func('mono.model.model_pipelines.' + model_type)(cfg)

    def forward(self, data):
        output = self.depth_model(**data)

        return output['prediction'], output['confidence'], output

    def inference(self, data):
        with torch.no_grad():
            pred_depth, confidence, _ = self.forward(data)
        return pred_depth, confidence
mono/model/model_pipelines/__base_model__.py
hugoycj/Metric3D_MVP

from .dense_pipeline import DensePredModel
from .__base_model__ import BaseDepthModel
__all__ = [
    'DensePredModel', 'BaseDepthModel',
]
mono/model/model_pipelines/__init__.py
hugoycj/Metric3D_MVP
import torch
import torch.nn as nn
from mono.utils.comm import get_func

class DensePredModel(nn.Module):
    def __init__(self, cfg) -> None:
        super(DensePredModel, self).__init__()

        self.encoder = get_func('mono.model.' + cfg.model.backbone.prefix + cfg.model.backbone.type)(**cfg.model.backbone)
        self.decoder = get_func('mono.model.' + cfg.model.decode_head.prefix + cfg.model.decode_head.type)(cfg)

    def forward(self, input, **kwargs):
        # [f_32, f_16, f_8, f_4]
        features = self.encoder(input)
        out = self.decoder(features, **kwargs)
        return out
mono/model/model_pipelines/dense_pipeline.py
hugoycj/Metric3D_MVP
import torch
import torch.nn as nn
from .model_pipelines.__base_model__ import BaseDepthModel

class DepthModel(BaseDepthModel):
    def __init__(self, cfg, **kwards):
        super(DepthModel, self).__init__(cfg)   
        model_type = cfg.model.type
        
    def inference(self, data):
        with torch.no_grad():
            pred_depth, confidence, output_dict = self.forward(data)       
        return pred_depth, confidence, output_dict

def get_monodepth_model(
    cfg : dict,
    **kwargs
    ) -> nn.Module:
    # config depth  model
    model = DepthModel(cfg, **kwargs)
    #model.init_weights(load_imagenet_model, imagenet_ckpt_fpath)
    assert isinstance(model, nn.Module)
    return model

def get_configured_monodepth_model(
    cfg: dict,
    ) -> nn.Module:
    """
        Args:
        @ configs: configures for the network.
        @ load_imagenet_model: whether to initialize from ImageNet-pretrained model.
        @ imagenet_ckpt_fpath: string representing path to file with weights to initialize model with.
        Returns:
        # model: depth model.
    """
    model = get_monodepth_model(cfg)
    return model

mono/model/monodepth_model.py
hugoycj/Metric3D_MVP
import os
import os.path as osp
import cv2
import time
import sys
CODE_SPACE=os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(CODE_SPACE)
import argparse
import mmcv
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

try:
    from mmcv.utils import Config, DictAction
except:
    from mmengine import Config, DictAction
from datetime import timedelta
import random
import numpy as np
from mono.utils.logger import setup_logger
import glob
from mono.utils.comm import init_env
from mono.model.monodepth_model import get_configured_monodepth_model
from mono.utils.running import load_ckpt
from mono.utils.do_test import do_scalecano_test_with_custom_data
from mono.utils.mldb import load_data_info, reset_ckpt_path
from mono.utils.custom_data import load_from_annos, load_data

def parse_args():
    parser = argparse.ArgumentParser(description='Train a segmentor')
    parser.add_argument('config', help='train config file path')
    parser.add_argument('--show-dir', help='the dir to save logs and visualization results')
    parser.add_argument('--load-from', help='the checkpoint file to load weights from')
    parser.add_argument('--node_rank', type=int, default=0)
    parser.add_argument('--nnodes', type=int, default=1, help='number of nodes')
    parser.add_argument('--options', nargs='+', action=DictAction, help='custom options')
    parser.add_argument('--launcher', choices=['None', 'pytorch', 'slurm', 'mpi', 'ror'], default='slurm', help='job launcher')
    parser.add_argument('--test_data_path', default='None', type=str, help='the path of test data')
    args = parser.parse_args()
    return args

def main(args):
    os.chdir(CODE_SPACE)
    cfg = Config.fromfile(args.config)
    
    if args.options is not None:
        cfg.merge_from_dict(args.options)
        
    # show_dir is determined in this priority: CLI > segment in file > filename
    if args.show_dir is not None:
        # update configs according to CLI args if args.show_dir is not None
        cfg.show_dir = args.show_dir
    else:
        # use condig filename + timestamp as default show_dir if args.show_dir is None
        cfg.show_dir = osp.join('./show_dirs', 
                                osp.splitext(osp.basename(args.config))[0],
                                args.timestamp)
    
    # ckpt path
    if args.load_from is None:
        raise RuntimeError('Please set model path!')
    cfg.load_from = args.load_from
    
    # load data info
    data_info = {}
    load_data_info('data_info', data_info=data_info)
    cfg.mldb_info = data_info
    # update check point info
    reset_ckpt_path(cfg.model, data_info)
    
    # create show dir
    os.makedirs(osp.abspath(cfg.show_dir), exist_ok=True)
    
    # init the logger before other steps
    cfg.log_file = osp.join(cfg.show_dir, f'{args.timestamp}.log')
    logger = setup_logger(cfg.log_file)
    
    # log some basic info
    logger.info(f'Config:\n{cfg.pretty_text}')
    
    # init distributed env dirst, since logger depends on the dist info
    if args.launcher == 'None':
        cfg.distributed = False
    else:
        cfg.distributed = True
        init_env(args.launcher, cfg)
    logger.info(f'Distributed training: {cfg.distributed}')
    
    # dump config 
    cfg.dump(osp.join(cfg.show_dir, osp.basename(args.config)))
    test_data_path = args.test_data_path
    if not os.path.isabs(test_data_path):
        test_data_path = osp.join(CODE_SPACE, test_data_path)

    if 'json' in test_data_path:
        test_data = load_from_annos(test_data_path)
    else:
        test_data = load_data(args.test_data_path)
    
    if not cfg.distributed:
        main_worker(0, cfg, args.launcher, test_data)
    else:
        # distributed training
        if args.launcher == 'ror':
            local_rank = cfg.dist_params.local_rank
            main_worker(local_rank, cfg, args.launcher, test_data)
        else:
            mp.spawn(main_worker, nprocs=cfg.dist_params.num_gpus_per_node, args=(cfg, args.launcher, test_data))
        
def main_worker(local_rank: int, cfg: dict, launcher: str, test_data: list):
    if cfg.distributed:
        cfg.dist_params.global_rank = cfg.dist_params.node_rank * cfg.dist_params.num_gpus_per_node + local_rank
        cfg.dist_params.local_rank = local_rank

        if launcher == 'ror':
            init_torch_process_group(use_hvd=False)
        else:
            torch.cuda.set_device(local_rank)
            default_timeout = timedelta(minutes=30)
            dist.init_process_group(
                backend=cfg.dist_params.backend,
                init_method=cfg.dist_params.dist_url,
                world_size=cfg.dist_params.world_size,
                rank=cfg.dist_params.global_rank,
                timeout=default_timeout)
    
    logger = setup_logger(cfg.log_file)
    # build model
    model = get_configured_monodepth_model(cfg, )
    
    # config distributed training
    if cfg.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model.cuda(),
                                                          device_ids=[local_rank],
                                                          output_device=local_rank,
                                                          find_unused_parameters=True)
    else:
        model = torch.nn.DataParallel(model).cuda()
        
    # load ckpt
    model, _,  _, _ = load_ckpt(cfg.load_from, model, strict_match=False)
    model.eval()
    
    do_scalecano_test_with_custom_data(
        model, 
        cfg,
        test_data,
        logger,
        cfg.distributed,
        local_rank
    )
    
if __name__ == '__main__':
    args = parse_args()
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    args.timestamp = timestamp
    main(args)    
mono/tools/test_scale_cano.py
hugoycj/Metric3D_MVP
import numpy as np
import torch
import torch.distributed as dist
import torch.nn.functional as F
import matplotlib.pyplot as plt


class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self) -> None:
        self.reset()

    def reset(self) -> None:
        self.val = np.longdouble(0.0)
        self.avg = np.longdouble(0.0)
        self.sum = np.longdouble(0.0)
        self.count = np.longdouble(0.0)

    def update(self, val, n: float = 1) -> None:
        self.val = val
        self.sum += val
        self.count += n
        self.avg = self.sum / (self.count + 1e-6)

class MetricAverageMeter(AverageMeter):
    """ 
    An AverageMeter designed specifically for evaluating segmentation results.
    """
    def __init__(self, metrics: list) -> None:
        """ Initialize object. """
        # average meters for metrics
        self.abs_rel = AverageMeter()
        self.rmse = AverageMeter()
        self.silog = AverageMeter()
        self.delta1 = AverageMeter()
        self.delta2 = AverageMeter()
        self.delta3 = AverageMeter()

        self.metrics = metrics

        self.log10 = AverageMeter()
        self.rmse_log = AverageMeter()
        self.sq_rel = AverageMeter()
    
    def reset(self):
        self.abs_rel.reset()
        self.rmse.reset()
        self.silog.reset()
        self.delta1.reset()
        self.delta2.reset()
        self.delta3.reset()
        self.log10.reset()
        self.rmse_log.reset()
        self.sq_rel.reset()

    def update_metrics_cpu(self,
        pred: torch.Tensor,
        target: torch.Tensor,
        mask: torch.Tensor,):
        '''
        Update metrics on cpu
        '''

        assert pred.shape == target.shape

        if len(pred.shape) == 3:
            pred = pred[:, None, :, :]
            target = target[:, None, :, :]
            mask = mask[:, None, :, :]
        elif len(pred.shape) == 2:
            pred = pred[None, None, :, :]
            target = target[None, None, :, :]
            mask = mask[None, None, :, :]
        

        # Absolute relative error
        abs_rel_sum, valid_pics = get_absrel_err(pred, target, mask)
        abs_rel_sum = abs_rel_sum.numpy()
        valid_pics = valid_pics.numpy()
        self.abs_rel.update(abs_rel_sum, valid_pics)

        # squraed relative error
        sqrel_sum, _ = get_sqrel_err(pred, target, mask)
        sqrel_sum = sqrel_sum.numpy()
        self.sq_rel.update(sqrel_sum, valid_pics)

        # root mean squared error
        rmse_sum, _ = get_rmse_err(pred, target, mask)
        rmse_sum = rmse_sum.numpy()
        self.rmse.update(rmse_sum, valid_pics)

        # log root mean squared error
        log_rmse_sum, _ = get_rmse_log_err(pred, target, mask)
        log_rmse_sum = log_rmse_sum.numpy()
        self.rmse_log.update(log_rmse_sum, valid_pics)

        # log10 error
        log10_sum, _ = get_log10_err(pred, target, mask)
        log10_sum = log10_sum.numpy()
        self.log10.update(log10_sum, valid_pics)

        # scale-invariant root mean squared error in log space
        silog_sum, _ = get_silog_err(pred, target, mask)
        silog_sum = silog_sum.numpy()
        self.silog.update(silog_sum, valid_pics)

        # ratio error, delta1, ...
        delta1_sum, delta2_sum, delta3_sum, _ = get_ratio_err(pred, target, mask)
        delta1_sum = delta1_sum.numpy()
        delta2_sum = delta2_sum.numpy()
        delta3_sum = delta3_sum.numpy()

        self.delta1.update(delta1_sum, valid_pics)
        self.delta2.update(delta2_sum, valid_pics)
        self.delta3.update(delta3_sum, valid_pics)


    def update_metrics_gpu(
        self,
        pred: torch.Tensor,
        target: torch.Tensor,
        mask: torch.Tensor,
        is_distributed: bool,
        pred_next: torch.Tensor = None,
        pose_f1_to_f2: torch.Tensor = None,
        intrinsic: torch.Tensor = None,
        ):
        '''
        Update metric on GPU. It supports distributed processing. If multiple machines are employed, please
        set 'is_distributed' as True.
        '''
        assert pred.shape == target.shape

        if len(pred.shape) == 3:
            pred = pred[:, None, :, :]
            target = target[:, None, :, :]
            mask = mask[:, None, :, :]
        elif len(pred.shape) == 2:
            pred = pred[None, None, :, :]
            target = target[None, None, :, :]
            mask = mask[None, None, :, :]
        

        # Absolute relative error
        abs_rel_sum, valid_pics = get_absrel_err(pred, target, mask)
        if is_distributed:
            dist.all_reduce(abs_rel_sum), dist.all_reduce(valid_pics)
        abs_rel_sum = abs_rel_sum.cpu().numpy()
        valid_pics = int(valid_pics)
        self.abs_rel.update(abs_rel_sum, valid_pics)

        # squraed relative error
        sqrel_sum, _ = get_sqrel_err(pred, target, mask)
        if is_distributed:
            dist.all_reduce(sqrel_sum)
        sqrel_sum = sqrel_sum.cpu().numpy()
        self.sq_rel.update(sqrel_sum, valid_pics)

        # root mean squared error
        rmse_sum, _ = get_rmse_err(pred, target, mask)
        if is_distributed:
            dist.all_reduce(rmse_sum)
        rmse_sum = rmse_sum.cpu().numpy()
        self.rmse.update(rmse_sum, valid_pics)

        # log root mean squared error
        log_rmse_sum, _ = get_rmse_log_err(pred, target, mask)
        if is_distributed:
            dist.all_reduce(log_rmse_sum)
        log_rmse_sum = log_rmse_sum.cpu().numpy()
        self.rmse_log.update(log_rmse_sum, valid_pics)

        # log10 error
        log10_sum, _ = get_log10_err(pred, target, mask)
        if is_distributed:
            dist.all_reduce(log10_sum)
        log10_sum = log10_sum.cpu().numpy()
        self.log10.update(log10_sum, valid_pics)

        # scale-invariant root mean squared error in log space
        silog_sum, _ = get_silog_err(pred, target, mask)
        if is_distributed:
            dist.all_reduce(silog_sum)
        silog_sum = silog_sum.cpu().numpy()
        self.silog.update(silog_sum, valid_pics)

        # ratio error, delta1, ...
        delta1_sum, delta2_sum, delta3_sum, _ = get_ratio_err(pred, target, mask)
        if is_distributed:
            dist.all_reduce(delta1_sum), dist.all_reduce(delta2_sum), dist.all_reduce(delta3_sum)
        delta1_sum = delta1_sum.cpu().numpy()
        delta2_sum = delta2_sum.cpu().numpy()
        delta3_sum = delta3_sum.cpu().numpy()

        self.delta1.update(delta1_sum, valid_pics)
        self.delta2.update(delta2_sum, valid_pics)
        self.delta3.update(delta3_sum, valid_pics)


    def get_metrics(self):
        metrics_dict = {}
        for metric in self.metrics:
            metrics_dict[metric] = self.__getattribute__(metric).avg
        return metrics_dict

def get_absrel_err(pred: torch.tensor,
                    target: torch.tensor,
                    mask: torch.tensor,
                    ):
    """
    Computes absolute relative error.
    Tasks preprocessed depths (no nans, infs and non-positive values).
    pred, target, and mask should be in the shape of [b, c, h, w]
    """

    assert len(pred.shape) == 4, len(target.shape) == 4
    b, c, h, w = pred.shape
    mask = mask.to(torch.float)
    t_m = target * mask
    p_m = pred * mask

    # Mean Absolute Relative Error
    rel = torch.abs(t_m - p_m) / (t_m + 1e-10) # compute errors
    abs_rel_sum = torch.sum(rel.reshape((b, c, -1)), dim=2) # [b, c]
    num = torch.sum(mask.reshape((b, c, -1)), dim=2) # [b, c]
    abs_err = abs_rel_sum / (num + 1e-10)
    valid_pics = torch.sum(num > 0)
    return torch.sum(abs_err), valid_pics

def get_sqrel_err(pred: torch.tensor,
                    target: torch.tensor,
                    mask: torch.tensor,
                    ):
    """
    Computes squared relative error.
    Tasks preprocessed depths (no nans, infs and non-positive values).
    pred, target, and mask should be in the shape of [b, c, h, w]
    """

    assert len(pred.shape) == 4, len(target.shape) == 4
    b, c, h, w = pred.shape
    mask = mask.to(torch.float)
    t_m = target * mask
    p_m = pred * mask

    # squared Relative Error
    sq_rel = torch.abs(t_m - p_m) ** 2 / (t_m + 1e-10) # compute errors
    sq_rel_sum = torch.sum(sq_rel.reshape((b, c, -1)), dim=2) # [b, c]
    num = torch.sum(mask.reshape((b, c, -1)), dim=2) # [b, c]
    sqrel_err = sq_rel_sum / (num + 1e-10)
    valid_pics = torch.sum(num > 0)
    return torch.sum(sqrel_err), valid_pics

def get_log10_err(pred: torch.tensor,
                    target: torch.tensor,
                    mask: torch.tensor,
                    ):
    """
    Computes log10 error.
    Tasks preprocessed depths (no nans, infs and non-positive values).
    pred, target, and mask should be in the shape of [b, c, h, w]
    """

    assert len(pred.shape) == 4, len(target.shape) == 4
    b, c, h, w = pred.shape
    mask = mask.to(torch.float)
    t_m = target * mask
    p_m = pred * mask

    diff_log = (torch.log10(p_m+1e-10) - torch.log10(t_m+1e-10)) * mask
    log10_diff = torch.abs(diff_log)
    log10_sum = torch.sum(log10_diff.reshape((b, c, -1)), dim=2) # [b, c]
    num = torch.sum(mask.reshape((b, c, -1)), dim=2) # [b, c]
    log10_err = log10_sum / (num + 1e-10)
    valid_pics = torch.sum(num > 0)
    return torch.sum(log10_err), valid_pics

def get_rmse_err(pred: torch.tensor,
                    target: torch.tensor,
                    mask: torch.tensor,
                    ):
    """
    Computes rmse error.
    Tasks preprocessed depths (no nans, infs and non-positive values).
    pred, target, and mask should be in the shape of [b, c, h, w]
    """

    assert len(pred.shape) == 4, len(target.shape) == 4
    b, c, h, w = pred.shape
    mask = mask.to(torch.float)
    t_m = target * mask
    p_m = pred * mask

    square = (t_m - p_m) ** 2
    rmse_sum = torch.sum(square.reshape((b, c, -1)), dim=2) # [b, c]
    num = torch.sum(mask.reshape((b, c, -1)), dim=2) # [b, c]
    rmse = torch.sqrt(rmse_sum / (num + 1e-10))
    valid_pics = torch.sum(num > 0)
    return torch.sum(rmse), valid_pics

def get_rmse_log_err(pred: torch.tensor,
                    target: torch.tensor,
                    mask: torch.tensor,
                    ):
    """
    Computes log rmse error.
    Tasks preprocessed depths (no nans, infs and non-positive values).
    pred, target, and mask should be in the shape of [b, c, h, w]
    """

    assert len(pred.shape) == 4, len(target.shape) == 4
    b, c, h, w = pred.shape
    mask = mask.to(torch.float)
    t_m = target * mask
    p_m = pred * mask

    diff_log = (torch.log10(p_m+1e-10) - torch.log10(t_m+1e-10)) * mask
    square = diff_log ** 2
    rmse_log_sum = torch.sum(square.reshape((b, c, -1)), dim=2) # [b, c]
    num = torch.sum(mask.reshape((b, c, -1)), dim=2) # [b, c]
    rmse_log = torch.sqrt(rmse_log_sum / (num + 1e-10))
    valid_pics = torch.sum(num > 0)
    return torch.sum(rmse_log), valid_pics

def get_silog_err(pred: torch.tensor,
                    target: torch.tensor,
                    mask: torch.tensor,
                    ):
    """
    Computes log rmse error.
    Tasks preprocessed depths (no nans, infs and non-positive values).
    pred, target, and mask should be in the shape of [b, c, h, w]
    """

    assert len(pred.shape) == 4, len(target.shape) == 4
    b, c, h, w = pred.shape
    mask = mask.to(torch.float)
    t_m = target * mask
    p_m = pred * mask

    diff_log = (torch.log10(p_m+1e-10) - torch.log10(t_m+1e-10)) * mask
    diff_log_sum = torch.sum(diff_log.reshape((b, c, -1)), dim=2) # [b, c]
    diff_log_square = diff_log ** 2
    diff_log_square_sum = torch.sum(diff_log_square.reshape((b, c, -1)), dim=2) # [b, c]
    num = torch.sum(mask.reshape((b, c, -1)), dim=2) # [b, c]
    silog = torch.sqrt(diff_log_square_sum / (num + 1e-10) - (diff_log_sum / (num + 1e-10)) ** 2)
    valid_pics = torch.sum(num > 0)
    return torch.sum(silog), valid_pics

def get_ratio_err(pred: torch.tensor,
                    target: torch.tensor,
                    mask: torch.tensor,
                    ):
    """
    Computes the percentage of pixels for which the ratio of the two depth maps is less than a given threshold.
    Tasks preprocessed depths (no nans, infs and non-positive values).
    pred, target, and mask should be in the shape of [b, c, h, w]
    """
    assert len(pred.shape) == 4, len(target.shape) == 4
    b, c, h, w = pred.shape
    mask = mask.to(torch.float)
    t_m = target * mask
    p_m = pred

    gt_pred = t_m / (p_m + 1e-10)
    pred_gt = p_m / (t_m + 1e-10)
    gt_pred = gt_pred.reshape((b, c, -1))
    pred_gt = pred_gt.reshape((b, c, -1))
    gt_pred_gt = torch.cat((gt_pred, pred_gt), axis=1)
    ratio_max = torch.amax(gt_pred_gt, axis=1)

    delta_1_sum = torch.sum((ratio_max < 1.25), dim=1) # [b, ]
    delta_2_sum = torch.sum((ratio_max < 1.25 ** 2), dim=1) # [b, ]
    delta_3_sum = torch.sum((ratio_max < 1.25 ** 3), dim=1) # [b, ]
    num = torch.sum(mask.reshape((b, -1)), dim=1) # [b, ]

    delta_1 = delta_1_sum / (num + 1e-10)
    delta_2 = delta_2_sum / (num + 1e-10)
    delta_3 = delta_3_sum / (num + 1e-10)
    valid_pics = torch.sum(num > 0)    

    return torch.sum(delta_1), torch.sum(delta_2), torch.sum(delta_3), valid_pics


if __name__ == '__main__':
    cfg = ['abs_rel', 'delta1']
    dam = MetricAverageMeter(cfg)

    pred_depth = np.random.random([2, 480, 640])
    gt_depth = np.random.random([2, 480, 640]) - 0.5
    intrinsic = [[100, 100, 200, 200], [200, 200, 300, 300]]
    
    pred = torch.from_numpy(pred_depth).cuda()
    gt = torch.from_numpy(gt_depth).cuda()

    mask = gt > 0
    dam.update_metrics_gpu(pred, gt, mask, False)
    eval_error = dam.get_metrics()
    print(eval_error)
    
mono/utils/avg_meter.py
hugoycj/Metric3D_MVP
import importlib 
import torch
import torch.distributed as dist
from .avg_meter import AverageMeter
from collections import defaultdict, OrderedDict
import os
import socket
from mmcv.utils import collect_env as collect_base_env
try:
    from mmcv.utils import get_git_hash
except:
    from mmengine.utils import get_git_hash
#import mono.mmseg as mmseg
# import mmseg
import time
import datetime
import logging


def main_process() -> bool:
    return get_rank() == 0
    #return not cfg.distributed or \
    #       (cfg.distributed and cfg.local_rank == 0)

def get_world_size() -> int:
    if not dist.is_available():
        return 1
    if not dist.is_initialized():
        return 1
    return dist.get_world_size()

def get_rank() -> int:
    if not dist.is_available():
        return 0
    if not dist.is_initialized():
        return 0
    return dist.get_rank()

def _find_free_port():
    # refer to https://github.com/facebookresearch/detectron2/blob/main/detectron2/engine/launch.py # noqa: E501
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # Binding to port 0 will cause the OS to find an available port for us
    sock.bind(('', 0))
    port = sock.getsockname()[1]
    sock.close()
    # NOTE: there is still a chance the port could be taken by other processes.
    return port 

def _is_free_port(port):
    ips = socket.gethostbyname_ex(socket.gethostname())[-1]
    ips.append('localhost')
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        return all(s.connect_ex((ip, port)) != 0 for ip in ips)


# def collect_env():
#     """Collect the information of the running environments."""
#     env_info = collect_base_env()
#     env_info['MMSegmentation'] = f'{mmseg.__version__}+{get_git_hash()[:7]}'

#     return env_info

def init_env(launcher, cfg):
    """Initialize distributed training environment.
    If argument ``cfg.dist_params.dist_url`` is specified as 'env://', then the master port will be system
    environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
    environment variable, then a default port ``29500`` will be used.
    """
    if launcher == 'slurm':
        _init_dist_slurm(cfg)
    elif launcher == 'ror':
        _init_dist_ror(cfg)
    elif launcher == 'None':
        _init_none_dist(cfg)
    else:
        raise RuntimeError(f'{cfg.launcher} has not been supported!')

def _init_none_dist(cfg):
    cfg.dist_params.num_gpus_per_node = 1
    cfg.dist_params.world_size = 1
    cfg.dist_params.nnodes = 1
    cfg.dist_params.node_rank = 0
    cfg.dist_params.global_rank = 0
    cfg.dist_params.local_rank = 0
    os.environ["WORLD_SIZE"] = str(1)

def _init_dist_ror(cfg):
    from ac2.ror.comm import get_local_rank, get_world_rank, get_local_size, get_node_rank, get_world_size
    cfg.dist_params.num_gpus_per_node = get_local_size()
    cfg.dist_params.world_size = get_world_size()
    cfg.dist_params.nnodes = (get_world_size()) // (get_local_size())
    cfg.dist_params.node_rank = get_node_rank()
    cfg.dist_params.global_rank = get_world_rank()
    cfg.dist_params.local_rank = get_local_rank()
    os.environ["WORLD_SIZE"] = str(get_world_size())


def _init_dist_slurm(cfg):
    if 'NNODES' not in os.environ:
        os.environ['NNODES'] = str(cfg.dist_params.nnodes)
    if 'NODE_RANK' not in os.environ:
        os.environ['NODE_RANK'] = str(cfg.dist_params.node_rank)

    #cfg.dist_params.
    num_gpus = torch.cuda.device_count()
    world_size = int(os.environ['NNODES']) * num_gpus
    os.environ['WORLD_SIZE'] = str(world_size)

    # config port
    if 'MASTER_PORT' in os.environ:
        master_port = str(os.environ['MASTER_PORT'])  # use MASTER_PORT in the environment variable
    else:
        # if torch.distributed default port(29500) is available
        # then use it, else find a free port
        if _is_free_port(16500):
            master_port = '16500'
        else:
            master_port = str(_find_free_port())
        os.environ['MASTER_PORT'] = master_port

    # config addr
    if 'MASTER_ADDR' in os.environ:
        master_addr = str(os.environ['MASTER_PORT'])  # use MASTER_PORT in the environment variable
    # elif cfg.dist_params.dist_url is not None:
    #     master_addr = ':'.join(str(cfg.dist_params.dist_url).split(':')[:2])
    else:
        master_addr = '127.0.0.1' #'tcp://127.0.0.1'
        os.environ['MASTER_ADDR'] = master_addr

    # set dist_url to 'env://' 
    cfg.dist_params.dist_url =  'env://' #f"{master_addr}:{master_port}"
    
    cfg.dist_params.num_gpus_per_node = num_gpus
    cfg.dist_params.world_size = world_size
    cfg.dist_params.nnodes = int(os.environ['NNODES'])
    cfg.dist_params.node_rank = int(os.environ['NODE_RANK'])
        
    # if int(os.environ['NNODES']) > 1 and cfg.dist_params.dist_url.startswith("file://"):
    #     raise Warning("file:// is not a reliable init_method in multi-machine jobs. Prefer tcp://")
        

def get_func(func_name):
    """
        Helper to return a function object by name. func_name must identify 
        a function in this module or the path to a function relative to the base
        module.
        @ func_name: function name.
    """
    if func_name == '':
        return None
    try:
        parts = func_name.split('.')
        # Refers to a function in this module
        if len(parts) == 1:
            return globals()[parts[0]]
        # Otherwise, assume we're referencing a module under modeling
        module_name = '.'.join(parts[:-1])
        module = importlib.import_module(module_name)
        return getattr(module, parts[-1])
    except:
        raise RuntimeError(f'Failed to find function: {func_name}')

class Timer(object):
    """A simple timer."""

    def __init__(self):
        self.reset()

    def tic(self):
        # using time.time instead of time.clock because time time.clock
        # does not normalize for multithreading
        self.start_time = time.time()

    def toc(self, average=True):
        self.diff = time.time() - self.start_time
        self.total_time += self.diff
        self.calls += 1
        self.average_time = self.total_time / self.calls
        if average:
            return self.average_time
        else:
            return self.diff

    def reset(self):
        self.total_time = 0.
        self.calls = 0
        self.start_time = 0.
        self.diff = 0.
        self.average_time = 0.

class TrainingStats(object):
    """Track vital training statistics."""
    def __init__(self, log_period, tensorboard_logger=None):
        self.log_period = log_period
        self.tblogger = tensorboard_logger
        self.tb_ignored_keys = ['iter', 'eta', 'epoch', 'time']
        self.iter_timer = Timer()
        # Window size for smoothing tracked values (with median filtering)
        self.filter_size = log_period
        def create_smoothed_value():
            return AverageMeter()
        self.smoothed_losses = defaultdict(create_smoothed_value)
        #self.smoothed_metrics = defaultdict(create_smoothed_value)
        #self.smoothed_total_loss = AverageMeter()


    def IterTic(self):
        self.iter_timer.tic()

    def IterToc(self):
        return self.iter_timer.toc(average=False)

    def reset_iter_time(self):
        self.iter_timer.reset()

    def update_iter_stats(self, losses_dict):
        """Update tracked iteration statistics."""
        for k, v in losses_dict.items():
            self.smoothed_losses[k].update(float(v), 1)

    def log_iter_stats(self, cur_iter, optimizer, max_iters, val_err={}):
        """Log the tracked statistics."""
        if (cur_iter % self.log_period == 0):
            stats = self.get_stats(cur_iter, optimizer, max_iters, val_err)
            log_stats(stats)
            if self.tblogger:
                self.tb_log_stats(stats, cur_iter)
            for k, v in self.smoothed_losses.items():
                v.reset()

    def tb_log_stats(self, stats, cur_iter):
        """Log the tracked statistics to tensorboard"""
        for k in stats:
            # ignore some logs
            if k not in self.tb_ignored_keys:
                v = stats[k]
                if isinstance(v, dict):
                    self.tb_log_stats(v, cur_iter)
                else:
                    self.tblogger.add_scalar(k, v, cur_iter)


    def get_stats(self, cur_iter, optimizer, max_iters, val_err = {}):
        eta_seconds = self.iter_timer.average_time * (max_iters - cur_iter)

        eta = str(datetime.timedelta(seconds=int(eta_seconds)))
        stats = OrderedDict(
            iter=cur_iter,  # 1-indexed
            time=self.iter_timer.average_time,
            eta=eta,
        )
        optimizer_state_dict = optimizer.state_dict()
        lr = {}
        for i in range(len(optimizer_state_dict['param_groups'])):
            lr_name = 'group%d_lr' % i
            lr[lr_name] = optimizer_state_dict['param_groups'][i]['lr']

        stats['lr'] = OrderedDict(lr)
        for k, v in self.smoothed_losses.items():
            stats[k] = v.avg

        stats['val_err'] = OrderedDict(val_err)
        stats['max_iters'] = max_iters
        return stats


def reduce_dict(input_dict, average=True):
    """
    Reduce the values in the dictionary from all processes so that process with rank
    0 has the reduced results.
    Args:
        @input_dict (dict): inputs to be reduced. All the values must be scalar CUDA Tensor.
        @average (bool): whether to do average or sum
    Returns:
        a dict with the same keys as input_dict, after reduction.
    """
    world_size = get_world_size()
    if world_size < 2:
        return input_dict
    with torch.no_grad():
        names = []
        values = []
        # sort the keys so that they are consistent across processes
        for k in sorted(input_dict.keys()):
            names.append(k)
            values.append(input_dict[k])
        values = torch.stack(values, dim=0)
        dist.reduce(values, dst=0)
        if dist.get_rank() == 0 and average:
            # only main process gets accumulated, so only divide by
            # world_size in this case
            values /= world_size
        reduced_dict = {k: v for k, v in zip(names, values)}
    return reduced_dict


def log_stats(stats):
    logger = logging.getLogger()
    """Log training statistics to terminal"""
    lines = "[Step %d/%d]\n" % (
            stats['iter'], stats['max_iters'])

    lines += "\t\tloss: %.3f,    time: %.6f,    eta: %s\n" % (
        stats['total_loss'], stats['time'], stats['eta'])

    # log loss
    lines += "\t\t" 
    for k, v in stats.items():
        if 'loss' in k.lower() and 'total_loss' not in k.lower():
            lines += "%s: %.3f" % (k, v)  + ",  "
    lines = lines[:-3]
    lines += '\n'

    # validate criteria
    lines += "\t\tlast val err:" + ",  ".join("%s: %.6f" % (k, v) for k, v in stats['val_err'].items()) + ", "
    lines += '\n'

    # lr in different groups
    lines += "\t\t" +  ",  ".join("%s: %.8f" % (k, v) for k, v in stats['lr'].items())
    lines += '\n'
    logger.info(lines[:-1])  # remove last new linen_pxl


mono/utils/comm.py
hugoycj/Metric3D_MVP
import glob
import os
import json
import cv2

def load_from_annos(anno_path):
    with open(anno_path, 'r') as f:
        annos = json.load(f)['files']

    datas = []
    for i, anno in enumerate(annos):
        rgb = anno['rgb']
        depth = anno['depth'] if 'depth' in anno else None
        depth_scale = anno['depth_scale'] if 'depth_scale' in anno else 1.0
        intrinsic = anno['cam_in'] if 'cam_in' in anno else None

        data_i = {
            'rgb': rgb,
            'depth': depth,
            'depth_scale': depth_scale,
            'intrinsic': intrinsic,
            'filename': os.path.basename(rgb),
            'folder': rgb.split('/')[-3],
        }
        datas.append(data_i)
    return datas

def load_data(path: str):
    rgbs = glob.glob(path + '/*.jpg') + glob.glob(path + '/*.png')
    data = [{'rgb':i, 'depth':None, 'intrinsic': None, 'filename':os.path.basename(i), 'folder': i.split('/')[-3]} for i in rgbs]
    return data
mono/utils/custom_data.py
hugoycj/Metric3D_MVP
import torch
import torch.nn.functional as F
import logging
import os
import os.path as osp
from mono.utils.avg_meter import MetricAverageMeter
from mono.utils.visualization import save_val_imgs, create_html, save_raw_imgs
import cv2
from tqdm import tqdm
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

from mono.utils.unproj_pcd import reconstruct_pcd, save_point_cloud

def to_cuda(data: dict):
    for k, v in data.items():
        if isinstance(v, torch.Tensor):
            data[k] = v.cuda(non_blocking=True)
        if isinstance(v, list) and len(v)>=1 and isinstance(v[0], torch.Tensor):
            for i, l_i in enumerate(v):
                data[k][i] = l_i.cuda(non_blocking=True)
    return data

def align_scale(pred: torch.tensor, target: torch.tensor):
    mask = target > 0
    if torch.sum(mask) > 10:
        scale = torch.median(target[mask]) / (torch.median(pred[mask]) + 1e-8)
    else:
        scale = 1
    pred_scaled = pred * scale
    return pred_scaled, scale

def align_scale_shift(pred: torch.tensor, target: torch.tensor):
    mask = target > 0
    target_mask = target[mask].cpu().numpy()
    pred_mask = pred[mask].cpu().numpy()
    if torch.sum(mask) > 10:
        scale, shift = np.polyfit(pred_mask, target_mask, deg=1)
        if scale < 0:
            scale = torch.median(target[mask]) / (torch.median(pred[mask]) + 1e-8)
            shift = 0
    else:
        scale = 1
        shift = 0
    pred = pred * scale + shift
    return pred, scale

def align_scale_shift_numpy(pred: np.array, target: np.array):
    mask = target > 0
    target_mask = target[mask]
    pred_mask = pred[mask]
    if np.sum(mask) > 10:
        scale, shift = np.polyfit(pred_mask, target_mask, deg=1)
        if scale < 0:
            scale = np.median(target[mask]) / (np.median(pred[mask]) + 1e-8)
            shift = 0
    else:
        scale = 1
        shift = 0
    pred = pred * scale + shift
    return pred, scale


def build_camera_model(H : int, W : int, intrinsics : list) -> np.array:
    """
    Encode the camera intrinsic parameters (focal length and principle point) to a 4-channel map. 
    """
    fx, fy, u0, v0 = intrinsics
    f = (fx + fy) / 2.0
    # principle point location
    x_row = np.arange(0, W).astype(np.float32)
    x_row_center_norm = (x_row - u0) / W
    x_center = np.tile(x_row_center_norm, (H, 1)) # [H, W]

    y_col = np.arange(0, H).astype(np.float32) 
    y_col_center_norm = (y_col - v0) / H
    y_center = np.tile(y_col_center_norm, (W, 1)).T # [H, W]

    # FoV
    fov_x = np.arctan(x_center / (f / W))
    fov_y = np.arctan(y_center / (f / H))

    cam_model = np.stack([x_center, y_center, fov_x, fov_y], axis=2)
    return cam_model

def resize_for_input(image, output_shape, intrinsic, canonical_shape, to_canonical_ratio):
    """
    Resize the input.
    Resizing consists of two processed, i.e. 1) to the canonical space (adjust the camera model); 2) resize the image while the camera model holds. Thus the
    label will be scaled with the resize factor.
    """
    padding = [123.675, 116.28, 103.53]
    h, w, _ = image.shape
    resize_ratio_h = output_shape[0] / canonical_shape[0]
    resize_ratio_w = output_shape[1] / canonical_shape[1]
    to_scale_ratio = min(resize_ratio_h, resize_ratio_w)

    resize_ratio = to_canonical_ratio * to_scale_ratio

    reshape_h = int(resize_ratio * h)
    reshape_w = int(resize_ratio * w)

    pad_h = max(output_shape[0] - reshape_h, 0)
    pad_w = max(output_shape[1] - reshape_w, 0)
    pad_h_half = int(pad_h / 2)
    pad_w_half = int(pad_w / 2)

    # resize
    image = cv2.resize(image, dsize=(reshape_w, reshape_h), interpolation=cv2.INTER_LINEAR)
    # padding
    image = cv2.copyMakeBorder(
        image, 
        pad_h_half, 
        pad_h - pad_h_half, 
        pad_w_half, 
        pad_w - pad_w_half, 
        cv2.BORDER_CONSTANT, 
        value=padding)
    
    # Resize, adjust principle point
    intrinsic[2] = intrinsic[2] * to_scale_ratio
    intrinsic[3] = intrinsic[3] * to_scale_ratio

    cam_model = build_camera_model(reshape_h, reshape_w, intrinsic)
    cam_model = cv2.copyMakeBorder(
        cam_model, 
        pad_h_half, 
        pad_h - pad_h_half, 
        pad_w_half, 
        pad_w - pad_w_half, 
        cv2.BORDER_CONSTANT, 
        value=-1)

    pad=[pad_h_half, pad_h - pad_h_half, pad_w_half, pad_w - pad_w_half]
    label_scale_factor=1/to_scale_ratio
    return image, cam_model, pad, label_scale_factor


def get_prediction(
    model: torch.nn.Module,
    input: torch.tensor,
    cam_model: torch.tensor,
    pad_info: torch.tensor,
    scale_info: torch.tensor,
    gt_depth: torch.tensor,
    normalize_scale: float,
    ori_shape: list=[],
):

    data = dict(
        input=input,
        cam_model=cam_model,
    )
    pred_depth, confidence, output_dict = model.module.inference(data)
    pred_depth = pred_depth.squeeze()
    pred_depth = pred_depth[pad_info[0] : pred_depth.shape[0] - pad_info[1], pad_info[2] : pred_depth.shape[1] - pad_info[3]]
    if gt_depth is not None:
        resize_shape = gt_depth.shape
    elif ori_shape != []:
        resize_shape = ori_shape
    else:
        resize_shape = pred_depth.shape

    pred_depth = torch.nn.functional.interpolate(pred_depth[None, None, :, :], resize_shape, mode='bilinear').squeeze() # to original size
    pred_depth = pred_depth * normalize_scale / scale_info
    if gt_depth is not None:
        pred_depth_scale, scale = align_scale(pred_depth, gt_depth)
    else:
        pred_depth_scale = None
        scale = None

    return pred_depth, pred_depth_scale, scale

def transform_test_data_scalecano(rgb, intrinsic, data_basic):
    """
    Pre-process the input for forwarding. Employ `label scale canonical transformation.'
        Args:
            rgb: input rgb image. [H, W, 3]
            intrinsic: camera intrinsic parameter, [fx, fy, u0, v0]
            data_basic: predefined canonical space in configs.
    """
    canonical_space = data_basic['canonical_space']
    forward_size = data_basic.crop_size
    mean = torch.tensor([123.675, 116.28, 103.53]).float()[:, None, None]
    std = torch.tensor([58.395, 57.12, 57.375]).float()[:, None, None]

    # BGR to RGB
    rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)

    ori_h, ori_w, _ = rgb.shape
    ori_focal = (intrinsic[0] + intrinsic[1]) / 2
    canonical_focal = canonical_space['focal_length']

    cano_label_scale_ratio = canonical_focal / ori_focal

    canonical_intrinsic = [
        intrinsic[0] * cano_label_scale_ratio,
        intrinsic[1] * cano_label_scale_ratio,
        intrinsic[2],
        intrinsic[3],
    ]

    # resize
    rgb, cam_model, pad, resize_label_scale_ratio = resize_for_input(rgb, forward_size, canonical_intrinsic, [ori_h, ori_w], 1.0)

    # label scale factor
    label_scale_factor = cano_label_scale_ratio * resize_label_scale_ratio

    rgb = torch.from_numpy(rgb.transpose((2, 0, 1))).float()
    rgb = torch.div((rgb - mean), std)
    rgb = rgb[None, :, :, :].cuda()
    
    cam_model = torch.from_numpy(cam_model.transpose((2, 0, 1))).float()
    cam_model = cam_model[None, :, :, :].cuda()
    cam_model_stacks = [
        torch.nn.functional.interpolate(cam_model, size=(cam_model.shape[2]//i, cam_model.shape[3]//i), mode='bilinear', align_corners=False)
        for i in [2, 4, 8, 16, 32]
    ]
    return rgb, cam_model_stacks, pad, label_scale_factor

def do_scalecano_test_with_custom_data(
    model: torch.nn.Module,
    cfg: dict,
    test_data: list,
    logger: logging.RootLogger,
    is_distributed: bool = True,
    local_rank: int = 0,
):

    show_dir = cfg.show_dir
    save_interval = 1
    save_imgs_dir = show_dir + '/vis'
    os.makedirs(save_imgs_dir, exist_ok=True)
    save_pcd_dir = show_dir + '/pcd'
    os.makedirs(save_pcd_dir, exist_ok=True)

    normalize_scale = cfg.data_basic.depth_range[1]
    dam = MetricAverageMeter(['abs_rel', 'rmse', 'silog', 'delta1', 'delta2', 'delta3'])
    dam_median = MetricAverageMeter(['abs_rel', 'rmse', 'silog', 'delta1', 'delta2', 'delta3'])
    dam_global = MetricAverageMeter(['abs_rel', 'rmse', 'silog', 'delta1', 'delta2', 'delta3'])
    
    for i, an in tqdm(enumerate(test_data)):
        rgb_origin = cv2.imread(an['rgb'])[:, :, ::-1].copy()
        if an['depth'] is not None:
            gt_depth = cv2.imread(an['depth'], -1)
            gt_depth_scale = an['depth_scale']
            gt_depth = gt_depth / gt_depth_scale
            gt_depth_flag = True
        else:
            gt_depth = None
            gt_depth_flag = False
        intrinsic = an['intrinsic']
        if intrinsic is None:
            intrinsic = [1000.0, 1000.0, rgb_origin.shape[1]/2, rgb_origin.shape[0]/2]
        rgb_input, cam_models_stacks, pad, label_scale_factor = transform_test_data_scalecano(rgb_origin, intrinsic, cfg.data_basic)

        pred_depth, pred_depth_scale, scale = get_prediction(
            model = model,
            input = rgb_input,
            cam_model = cam_models_stacks,
            pad_info = pad,
            scale_info = label_scale_factor,
            gt_depth = None,
            normalize_scale = normalize_scale,
            ori_shape=[rgb_origin.shape[0], rgb_origin.shape[1]],
        )

        if gt_depth_flag:

            pred_depth = torch.nn.functional.interpolate(pred_depth[None, None, :, :], (gt_depth.shape[0], gt_depth.shape[1]), mode='bilinear').squeeze() # to original size
            gt_depth = torch.from_numpy(gt_depth).cuda()

            pred_depth_median = pred_depth * gt_depth[gt_depth != 0].median() / pred_depth[gt_depth != 0].median()
            pred_global, _ = align_scale_shift(pred_depth, gt_depth)
            
            mask = (gt_depth > 1e-8)
            dam.update_metrics_gpu(pred_depth, gt_depth, mask, is_distributed)
            dam_median.update_metrics_gpu(pred_depth_median, gt_depth, mask, is_distributed)
            dam_global.update_metrics_gpu(pred_global, gt_depth, mask, is_distributed)
        
        if i % save_interval == 0:
            os.makedirs(osp.join(save_imgs_dir, an['folder']), exist_ok=True)
            rgb_torch = torch.from_numpy(rgb_origin).to(pred_depth.device).permute(2, 0, 1)
            mean = torch.tensor([123.675, 116.28, 103.53]).float()[:, None, None].to(rgb_torch.device)
            std = torch.tensor([58.395, 57.12, 57.375]).float()[:, None, None].to(rgb_torch.device)
            rgb_torch = torch.div((rgb_torch - mean), std)

            save_val_imgs(
                i,
                pred_depth,
                gt_depth if gt_depth is not None else torch.ones_like(pred_depth, device=pred_depth.device),
                rgb_torch,
                osp.join(an['folder'], an['filename']),
                save_imgs_dir,
            )

            # pcd
            pred_depth = pred_depth.detach().cpu().numpy()
            pcd = reconstruct_pcd(pred_depth, intrinsic[0], intrinsic[1], intrinsic[2], intrinsic[3])
            os.makedirs(osp.join(save_pcd_dir, an['folder']), exist_ok=True)
            
            if an['intrinsic'] != None:
                save_point_cloud(pcd.reshape((-1, 3)), rgb_origin.reshape(-1, 3), osp.join(save_pcd_dir, an['folder'], an['filename'][:-4]+'.ply'))
            else:
                # for r in [0.9, 1.0, 1.1]:
                #     for f in [600, 800, 1000, 1250, 1500]:
                for r in [1.0]:
                    for f in [1000]:
                        #f_scale = f
                        f_scale = f * (rgb_origin.shape[0] + rgb_origin.shape[1]) / (cfg.data_basic.canonical_space['img_size'][0] + cfg.data_basic.canonical_space['img_size'][1])
                        pcd = reconstruct_pcd(pred_depth, f_scale * r, f_scale * (2 - r), intrinsic[2], intrinsic[3])
                        fstr = '_fx_' + str(int(f_scale * r)) + '_fy_' + str(int(f_scale * (2-r)))
                        save_point_cloud(pcd.reshape((-1, 3)), rgb_origin.reshape(-1, 3), osp.join(save_pcd_dir, an['folder'], an['filename'][:-4]+fstr+'.ply'))
    
    if gt_depth_flag:
        eval_error = dam.get_metrics()
        print('w/o match :', eval_error)

        eval_error_median = dam_median.get_metrics()
        print('median match :', eval_error_median)

        eval_error_global = dam_global.get_metrics()
        print('global match :', eval_error_global)
    else:
        print('missing gt_depth, only save visualizations...')

mono/utils/do_test.py
hugoycj/Metric3D_MVP
import atexit
import logging
import os
import sys
import time
import torch
from termcolor import colored

__all__ = ["setup_logger", ]

class _ColorfulFormatter(logging.Formatter):
    def __init__(self, *args, **kwargs):
        self._root_name = kwargs.pop("root_name") + "."
        self._abbrev_name = kwargs.pop("abbrev_name", "")
        if len(self._abbrev_name):
            self._abbrev_name = self._abbrev_name + "."
        super(_ColorfulFormatter, self).__init__(*args, **kwargs)

    def formatMessage(self, record):
        record.name = record.name.replace(self._root_name, self._abbrev_name)
        log = super(_ColorfulFormatter, self).formatMessage(record)
        if record.levelno == logging.WARNING:
            prefix = colored("WARNING", "red", attrs=["blink"])
        elif record.levelno == logging.ERROR or record.levelno == logging.CRITICAL:
            prefix = colored("ERROR", "red", attrs=["blink", "underline"])
        else:
            return log
        return prefix + " " + log

def setup_logger(
    output=None, distributed_rank=0, *, name='metricdepth', color=True, abbrev_name=None
):
    """
    Initialize the detectron2 logger and set its verbosity level to "DEBUG".
    Args:
        output (str): a file name or a directory to save log. If None, will not save log file.
            If ends with ".txt" or ".log", assumed to be a file name.
            Otherwise, logs will be saved to `output/log.txt`.
        abbrev_name (str): an abbreviation of the module, to avoid log names in logs.
            Set to "" not log the root module in logs.
            By default, will abbreviate "detectron2" to "d2" and leave other
            modules unchanged.
    Returns:
        logging.Logger: a logger
    """
    logger = logging.getLogger()
    logger.setLevel(logging.INFO) # NOTE: if more detailed, change it to logging.DEBUG
    logger.propagate = False

    if abbrev_name is None:
        abbrev_name = "d2"
        
    plain_formatter = logging.Formatter(
        "[%(asctime)s] %(name)s %(levelname)s %(message)s ", datefmt="%m/%d %H:%M:%S"
    )
    # stdout logging: master  only
    if distributed_rank == 0:
        ch = logging.StreamHandler(stream=sys.stdout)
        ch.setLevel(logging.INFO) # NOTE: if more detailed, change it to logging.DEBUG
        if color:
            formatter = _ColorfulFormatter(
                colored("[%(asctime)s %(name)s]: ", "green") + "%(message)s",
                datefmt="%m/%d %H:%M:%S",
                root_name=name,
                abbrev_name=str(abbrev_name),
            )
        else:
            formatter = plain_formatter
        ch.setFormatter(formatter)
        logger.addHandler(ch)
    
    # file logging: all workers
    if output is not None:
        if output.endswith(".txt") or output.endswith(".log"):
            filename = output
        else:
            filename = os.path.join(output, "log.txt")
        if distributed_rank > 0:
            filename = filename + ".rank{}".format(distributed_rank)
        os.makedirs(os.path.dirname(filename), exist_ok=True)

        fh = logging.StreamHandler(_cached_log_stream(filename))
        fh.setLevel(logging.INFO) # NOTE: if more detailed, change it to logging.DEBUG
        fh.setFormatter(plain_formatter)
        logger.addHandler(fh)
    

    return logger

from iopath.common.file_io import PathManager as PathManagerBase


PathManager = PathManagerBase()

# cache the opened file object, so that different calls to 'setup_logger
# with the same file name can safely write to the same file.
def _cached_log_stream(filename):
    # use 1K buffer if writting to cloud storage
    io = PathManager.open(filename, "a", buffering=1024 if "://" in filename else -1)
    atexit.register(io.close)
    return io    
    
mono/utils/logger.py
hugoycj/Metric3D_MVP
from types import ModuleType
import data_info

def load_data_info(module_name, data_info={}, mldb_type='mldb_info', module=None):
    if module is None:
        module = globals().get(module_name, None)
    if module:
        for key, value in module.__dict__.items():
            if not (key.startswith('__')) and not (key.startswith('_')):
                if key == 'mldb_info':
                    data_info.update(value)
                elif isinstance(value, ModuleType):
                    load_data_info(module_name + '.' + key, data_info, module=value)
    else:
        raise RuntimeError(f'Try to access "mldb_info", but cannot find {module_name} module.')

def reset_ckpt_path(cfg, data_info):
    if isinstance(cfg, dict):
        for key in cfg.keys():
            if key == 'backbone':
                new_ckpt_path = data_info['checkpoint']['mldb_root'] + '/' + data_info['checkpoint'][cfg.backbone.type]
                cfg.backbone.update(checkpoint=new_ckpt_path)
                continue
            elif isinstance(cfg.get(key), dict):
                reset_ckpt_path(cfg.get(key), data_info)
            else:
                continue
    else:
        return

if __name__ == '__main__':
    mldb_info_tmp = {}
    load_data_info('mldb_data_info', mldb_info_tmp)
    print('results', mldb_info_tmp.keys())
mono/utils/mldb.py
hugoycj/Metric3D_MVP
import os
import torch
import torch.nn as nn
from mono.utils.comm import main_process
import copy
import inspect
import logging
import glob


def load_ckpt(load_path, model, optimizer=None, scheduler=None, strict_match=True, loss_scaler=None):
    """
    Load the check point for resuming training or finetuning.
    """
    logger = logging.getLogger()
    if os.path.isfile(load_path):
        if main_process():
            logger.info(f"Loading weight '{load_path}'")
        checkpoint = torch.load(load_path, map_location="cpu")
        ckpt_state_dict  = checkpoint['model_state_dict']
        model.module.load_state_dict(ckpt_state_dict, strict=strict_match)

        if optimizer is not None:
            optimizer.load_state_dict(checkpoint['optimizer'])
        if scheduler is not None:
            scheduler.load_state_dict(checkpoint['scheduler'])
        if loss_scaler is not None and 'scaler' in checkpoint:
            scheduler.load_state_dict(checkpoint['scaler'])
        del ckpt_state_dict
        del checkpoint
        if main_process():
            logger.info(f"Successfully loaded weight: '{load_path}'")
            if scheduler is not None and optimizer is not None:
                logger.info(f"Resume training from: '{load_path}'")
    else:
        if main_process():
            raise RuntimeError(f"No weight found at '{load_path}'")
    return model, optimizer, scheduler, loss_scaler


def save_ckpt(cfg, model, optimizer, scheduler, curr_iter=0, curr_epoch=None, loss_scaler=None):
    """
    Save the model, optimizer, lr scheduler.
    """
    logger = logging.getLogger()

    if 'IterBasedRunner' in cfg.runner.type:
        max_iters = cfg.runner.max_iters
    elif 'EpochBasedRunner' in cfg.runner.type:
        max_iters = cfg.runner.max_epochs
    else:
        raise TypeError(f'{cfg.runner.type} is not supported')

    ckpt = dict(
        model_state_dict=model.module.state_dict(),
        optimizer=optimizer.state_dict(),
        max_iter=cfg.runner.max_iters if 'max_iters' in cfg.runner \
            else cfg.runner.max_epochs,
        scheduler=scheduler.state_dict(),
    )

    if loss_scaler is not None:
        ckpt.update(dict(scaler=loss_scaler.state_dict()))
    
    ckpt_dir = os.path.join(cfg.work_dir, 'ckpt')
    os.makedirs(ckpt_dir, exist_ok=True)

    save_name = os.path.join(ckpt_dir, 'step%08d.pth' %curr_iter)
    saved_ckpts = glob.glob(ckpt_dir + '/step*.pth')
    torch.save(ckpt, save_name)

    # keep the last 8 ckpts
    if len(saved_ckpts) > 20:
        saved_ckpts.sort()
        os.remove(saved_ckpts.pop(0))
    
    logger.info(f'Save model: {save_name}')

mono/utils/running.py
hugoycj/Metric3D_MVP
import collections
import cv2
import math
import numpy as np
import numbers
import random
import torch

import matplotlib
import matplotlib.cm


"""
Provides a set of Pytorch transforms that use OpenCV instead of PIL (Pytorch default)
for image manipulation.
"""

class Compose(object):
    # Composes transforms: transforms.Compose([transforms.RandScale([0.5, 2.0]), transforms.ToTensor()])
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, images, labels, intrinsics, cam_models=None, other_labels=None, transform_paras=None):
        for t in self.transforms:
            images, labels, intrinsics, cam_models, other_labels, transform_paras = t(images, labels, intrinsics, cam_models, other_labels, transform_paras)
        return images, labels, intrinsics, cam_models, other_labels, transform_paras


class ToTensor(object):
    # Converts numpy.ndarray (H x W x C) to a torch.FloatTensor of shape (C x H x W).
    def __init__(self,  **kwargs):
        return
    def __call__(self, images, labels, intrinsics, cam_models=None, other_labels=None, transform_paras=None):
        if not isinstance(images, list) or not isinstance(labels, list) or not isinstance(intrinsics, list):
            raise (RuntimeError("transform.ToTensor() only handle inputs/labels/intrinsics lists."))
        if len(images) != len(intrinsics):
            raise (RuntimeError("Numbers of images and intrinsics are not matched."))
        if not isinstance(images[0], np.ndarray) or not isinstance(labels[0], np.ndarray):
            raise (RuntimeError("transform.ToTensor() only handle np.ndarray for the input and label."
                                "[eg: data readed by cv2.imread()].\n"))
        if  not isinstance(intrinsics[0], list):
            raise (RuntimeError("transform.ToTensor() only handle list for the camera intrinsics"))

        if len(images[0].shape) > 3 or len(images[0].shape) < 2:
            raise (RuntimeError("transform.ToTensor() only handle image(np.ndarray) with 3 dims or 2 dims.\n"))
        if len(labels[0].shape) > 3 or len(labels[0].shape) < 2:
            raise (RuntimeError("transform.ToTensor() only handle label(np.ndarray) with 3 dims or 2 dims.\n"))

        if len(intrinsics[0]) >4 or len(intrinsics[0]) < 3:
            raise (RuntimeError("transform.ToTensor() only handle intrinsic(list) with 3 sizes or 4 sizes.\n"))
        
        for i, img in enumerate(images):
            if len(img.shape) == 2:
                img = np.expand_dims(img, axis=2)
            images[i] = torch.from_numpy(img.transpose((2, 0, 1))).float()
        for i, lab in enumerate(labels):
            if len(lab.shape) == 2:
                lab = np.expand_dims(lab, axis=0)
            labels[i] = torch.from_numpy(lab).float()
        for i, intrinsic in enumerate(intrinsics):
            if len(intrinsic) == 3:
                intrinsic = [intrinsic[0],] + intrinsic
            intrinsics[i] = torch.tensor(intrinsic, dtype=torch.float)
        if cam_models is not None:
            for i, cam_model in enumerate(cam_models):
                cam_models[i] = torch.from_numpy(cam_model.transpose((2, 0, 1))).float() if cam_model is not None else None
        if other_labels is not None:
            for i, lab in enumerate(other_labels):
                if len(lab.shape) == 2:
                    lab = np.expand_dims(lab, axis=0)
                other_labels[i] = torch.from_numpy(lab).float()
        return images, labels, intrinsics, cam_models, other_labels, transform_paras


class Normalize(object):
    # Normalize tensor with mean and standard deviation along channel: channel = (channel - mean) / std
    def __init__(self, mean, std=None, **kwargs):
        if std is None:
            assert len(mean) > 0
        else:
            assert len(mean) == len(std)
        self.mean = torch.tensor(mean).float()[:, None, None]
        self.std = torch.tensor(std).float()[:, None, None] if std is not None \
            else torch.tensor([1.0, 1.0, 1.0]).float()[:, None, None]

    def __call__(self, images, labels, intrinsics, cam_models=None, other_labels=None, transform_paras=None):
        # if self.std is None:
        #     # for t, m in zip(image, self.mean):
        #     #     t.sub(m)
        #     image = image - self.mean
        #     if ref_images is not None:
        #         for i, ref_i in enumerate(ref_images):
        #             ref_images[i] =  ref_i - self.mean
        # else:
        #     # for t, m, s in zip(image, self.mean, self.std):
        #     #     t.sub(m).div(s)
        #     image = (image - self.mean) / self.std
        #     if ref_images is not None:
        #         for i, ref_i in enumerate(ref_images):
        #             ref_images[i] =  (ref_i - self.mean) / self.std
        for i, img in enumerate(images):
            img = torch.div((img - self.mean), self.std)
            images[i] = img
        return images, labels, intrinsics, cam_models, other_labels, transform_paras


class LableScaleCanonical(object):
    """
    To solve the ambiguity observation for the mono branch, i.e. different focal length (object size) with the same depth, cameras are
    mapped to a canonical space. To mimic this, we set the focal length to a canonical one and scale the depth value. NOTE: resize the image based on the ratio can also solve
    Args:
        images: list of RGB images.
        labels: list of depth/disparity labels.
        other labels: other labels, such as instance segmentations, semantic segmentations...
    """
    def __init__(self, **kwargs):
        self.canonical_focal = kwargs['focal_length']
    
    def _get_scale_ratio(self, intrinsic):
        target_focal_x = intrinsic[0]
        label_scale_ratio = self.canonical_focal / target_focal_x
        pose_scale_ratio = 1.0
        return label_scale_ratio, pose_scale_ratio
    
    def __call__(self, images, labels, intrinsics, cam_models=None, other_labels=None, transform_paras=None):
        assert len(images[0].shape) == 3 and len(labels[0].shape) == 2
        assert labels[0].dtype == np.float32
        
        label_scale_ratio = None
        pose_scale_ratio = None

        for i in range(len(intrinsics)):
            img_i = images[i]
            label_i = labels[i] if i < len(labels) else None
            intrinsic_i = intrinsics[i].copy()
            cam_model_i = cam_models[i] if cam_models is not None and i < len(cam_models) else None

            label_scale_ratio, pose_scale_ratio = self._get_scale_ratio(intrinsic_i)

            # adjust the focal length, map the current camera to the canonical space
            intrinsics[i] = [intrinsic_i[0] * label_scale_ratio, intrinsic_i[1] * label_scale_ratio, intrinsic_i[2], intrinsic_i[3]]

            # scale the label to the canonical space
            if label_i is not None:
                labels[i] = label_i * label_scale_ratio
            
            if cam_model_i is not None:
                # As the focal length is adjusted (canonical focal length), the camera model should be re-built
                ori_h, ori_w, _ = img_i.shape
                cam_models[i] = build_camera_model(ori_h, ori_w, intrinsics[i])
            

        if transform_paras is not None:
            transform_paras.update(label_scale_factor=label_scale_ratio, focal_scale_factor=label_scale_ratio)
        
        return images, labels, intrinsics, cam_models, other_labels, transform_paras


class ResizeKeepRatio(object):
    """
    Resize and pad to a given size. Hold the aspect ratio.
    This resizing assumes that the camera model remains unchanged.
    Args:
        resize_size: predefined output size.
    """
    def __init__(self, resize_size, padding=None, ignore_label=-1, **kwargs):
        if isinstance(resize_size, int):
            self.resize_h = resize_size
            self.resize_w = resize_size
        elif isinstance(resize_size, collections.Iterable) and len(resize_size) == 2 \
                and isinstance(resize_size[0], int) and isinstance(resize_size[1], int) \
                and resize_size[0] > 0 and resize_size[1] > 0:
            self.resize_h = resize_size[0]
            self.resize_w = resize_size[1]
        else:
            raise (RuntimeError("crop size error.\n"))
        if padding is None:
            self.padding = padding
        elif isinstance(padding, list):
            if all(isinstance(i, numbers.Number) for i in padding):
                self.padding = padding
            else:
                raise (RuntimeError("padding in Crop() should be a number list\n"))
            if len(padding) != 3:
                raise (RuntimeError("padding channel is not equal with 3\n"))
        else:
            raise (RuntimeError("padding in Crop() should be a number list\n"))
        if isinstance(ignore_label, int):
            self.ignore_label = ignore_label
        else:
            raise (RuntimeError("ignore_label should be an integer number\n"))
        # self.crop_size = kwargs['crop_size']
        self.canonical_focal = kwargs['focal_length']
        
    def main_data_transform(self, image, label, intrinsic, cam_model, resize_ratio, padding, to_scale_ratio):
        """
        Resize data first and then do the padding.
        'label' will be scaled.
        """
        h, w, _ = image.shape
        reshape_h = int(resize_ratio * h)
        reshape_w = int(resize_ratio * w)

        pad_h, pad_w, pad_h_half, pad_w_half = padding
        
        # resize
        image = cv2.resize(image, dsize=(reshape_w, reshape_h), interpolation=cv2.INTER_LINEAR)
        # padding
        image = cv2.copyMakeBorder(
            image, 
            pad_h_half, 
            pad_h - pad_h_half, 
            pad_w_half, 
            pad_w - pad_w_half, 
            cv2.BORDER_CONSTANT, 
            value=self.padding)

        if label is not None:
            # label = cv2.resize(label, dsize=(reshape_w, reshape_h), interpolation=cv2.INTER_NEAREST)
            label = resize_depth_preserve(label, (reshape_h, reshape_w))
            label = cv2.copyMakeBorder(
                label, 
                pad_h_half, 
                pad_h - pad_h_half, 
                pad_w_half, 
                pad_w - pad_w_half, 
                cv2.BORDER_CONSTANT, 
                value=self.ignore_label)
            # scale the label
            label = label / to_scale_ratio
        
        # Resize, adjust principle point
        if intrinsic is not None:
            intrinsic[0] = intrinsic[0] * resize_ratio / to_scale_ratio
            intrinsic[1] = intrinsic[1] * resize_ratio / to_scale_ratio
            intrinsic[2] = intrinsic[2] * resize_ratio
            intrinsic[3] = intrinsic[3] * resize_ratio

        if cam_model is not None:
            #cam_model = cv2.resize(cam_model, dsize=(reshape_w, reshape_h), interpolation=cv2.INTER_LINEAR)
            cam_model = build_camera_model(reshape_h, reshape_w, intrinsic)
            cam_model = cv2.copyMakeBorder(
                cam_model, 
                pad_h_half, 
                pad_h - pad_h_half, 
                pad_w_half, 
                pad_w - pad_w_half, 
                cv2.BORDER_CONSTANT, 
                value=self.ignore_label)

        # Pad, adjust the principle point
        if intrinsic is not None:
            intrinsic[2] = intrinsic[2] + pad_w_half
            intrinsic[3] = intrinsic[3] + pad_h_half
        return image, label, intrinsic, cam_model

    def get_label_scale_factor(self, image, intrinsic, resize_ratio):
        ori_h, ori_w, _ = image.shape
        # crop_h, crop_w = self.crop_size
        ori_focal = intrinsic[0]

        to_canonical_ratio = self.canonical_focal / ori_focal
        to_scale_ratio = resize_ratio / to_canonical_ratio
        return to_scale_ratio

    def __call__(self, images, labels, intrinsics, cam_models=None, other_labels=None, transform_paras=None):
        target_h, target_w, _ = images[0].shape
        resize_ratio_h = self.resize_h / target_h
        resize_ratio_w = self.resize_w / target_w
        resize_ratio = min(resize_ratio_h, resize_ratio_w)
        reshape_h = int(resize_ratio * target_h)
        reshape_w = int(resize_ratio * target_w)
        pad_h = max(self.resize_h - reshape_h, 0)
        pad_w = max(self.resize_w - reshape_w, 0)
        pad_h_half = int(pad_h / 2)
        pad_w_half = int(pad_w / 2)

        pad_info = [pad_h, pad_w, pad_h_half, pad_w_half]
        to_scale_ratio = self.get_label_scale_factor(images[0], intrinsics[0], resize_ratio)

        for i in range(len(images)):
            img = images[i]
            label = labels[i] if i < len(labels) else None
            intrinsic = intrinsics[i] if i < len(intrinsics) else None
            cam_model = cam_models[i] if cam_models is not None and i < len(cam_models) else None
            img, label, intrinsic, cam_model = self.main_data_transform(
                img, label, intrinsic, cam_model, resize_ratio, pad_info, to_scale_ratio)
            images[i] = img
            if label is not None:
                labels[i] = label
            if intrinsic is not None:
                intrinsics[i] = intrinsic
            if cam_model is not None:
                cam_models[i] = cam_model
        
        if other_labels is not None:
            
            for i, other_lab in enumerate(other_labels):
                # resize
                other_lab =  cv2.resize(other_lab, dsize=(reshape_w, reshape_h), interpolation=cv2.INTER_NEAREST)
                # pad
                other_labels[i] =  cv2.copyMakeBorder(
                    other_lab, 
                    pad_h_half, 
                    pad_h - pad_h_half, 
                    pad_w_half, 
                    pad_w - pad_w_half, 
                    cv2.BORDER_CONSTANT, 
                    value=self.ignore_label)

        pad = [pad_h_half, pad_h - pad_h_half, pad_w_half, pad_w - pad_w_half]
        if transform_paras is not None:
            pad_old = transform_paras['pad'] if 'pad' in transform_paras else [0,0,0,0]
            new_pad = [pad_old[0] + pad[0], pad_old[1] + pad[1], pad_old[2] + pad[2], pad_old[3] + pad[3]]
            transform_paras.update(dict(pad=new_pad))
            if 'label_scale_factor' in transform_paras:
                transform_paras['label_scale_factor'] = transform_paras['label_scale_factor'] * 1.0 / to_scale_ratio
            else:
                transform_paras.update(label_scale_factor=1.0/to_scale_ratio)
        return images, labels, intrinsics, cam_models, other_labels, transform_paras


class BGR2RGB(object):
    # Converts image from BGR order to RGB order, for model initialized from Pytorch
    def __init__(self,  **kwargs):
        return
    def __call__(self, images, labels, intrinsics, cam_models=None,other_labels=None, transform_paras=None):
        for i, img in enumerate(images):
            images[i] = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        return images, labels, intrinsics, cam_models, other_labels, transform_paras
    
    
def resize_depth_preserve(depth, shape):
    """
    Resizes depth map preserving all valid depth pixels
    Multiple downsampled points can be assigned to the same pixel.

    Parameters
    ----------
    depth : np.array [h,w]
        Depth map
    shape : tuple (H,W)
        Output shape

    Returns
    -------
    depth : np.array [H,W,1]
        Resized depth map
    """
    # Store dimensions and reshapes to single column
    depth = np.squeeze(depth)
    h, w = depth.shape
    x = depth.reshape(-1)
    # Create coordinate grid
    uv = np.mgrid[:h, :w].transpose(1, 2, 0).reshape(-1, 2)
    # Filters valid points
    idx = x > 0
    crd, val = uv[idx], x[idx]
    # Downsamples coordinates
    crd[:, 0] = (crd[:, 0] * (shape[0] / h) + 0.5).astype(np.int32)
    crd[:, 1] = (crd[:, 1] * (shape[1] / w) + 0.5).astype(np.int32)
    # Filters points inside image
    idx = (crd[:, 0] < shape[0]) & (crd[:, 1] < shape[1])
    crd, val = crd[idx], val[idx]
    # Creates downsampled depth image and assigns points
    depth = np.zeros(shape)
    depth[crd[:, 0], crd[:, 1]] = val
    # Return resized depth map
    return depth


def build_camera_model(H : int, W : int, intrinsics : list) -> np.array:
    """
    Encode the camera intrinsic parameters (focal length and principle point) to a 4-channel map. 
    """
    fx, fy, u0, v0 = intrinsics
    f = (fx + fy) / 2.0
    # principle point location
    x_row = np.arange(0, W).astype(np.float32)
    x_row_center_norm = (x_row - u0) / W
    x_center = np.tile(x_row_center_norm, (H, 1)) # [H, W]

    y_col = np.arange(0, H).astype(np.float32) 
    y_col_center_norm = (y_col - v0) / H
    y_center = np.tile(y_col_center_norm, (W, 1)).T

    # FoV
    fov_x = np.arctan(x_center / (f / W))
    fov_y =  np.arctan(y_center/ (f / H))

    cam_model = np.stack([x_center, y_center, fov_x, fov_y], axis=2)
    return cam_model

def gray_to_colormap(img, cmap='rainbow'):
    """
    Transfer gray map to matplotlib colormap
    """
    assert img.ndim == 2

    img[img<0] = 0
    mask_invalid = img < 1e-10
    img = img / (img.max() + 1e-8)
    norm = matplotlib.colors.Normalize(vmin=0, vmax=1.1)
    cmap_m = matplotlib.cm.get_cmap(cmap)
    map = matplotlib.cm.ScalarMappable(norm=norm, cmap=cmap_m)
    colormap = (map.to_rgba(img)[:, :, :3] * 255).astype(np.uint8)
    colormap[mask_invalid] = 0
    return colormap
mono/utils/transform.py
hugoycj/Metric3D_MVP
import numpy as np
import torch
from plyfile import PlyData, PlyElement
import cv2


def get_pcd_base(H, W, u0, v0, fx, fy):
    x_row = np.arange(0, W)
    x = np.tile(x_row, (H, 1))
    x = x.astype(np.float32)
    u_m_u0 = x - u0

    y_col = np.arange(0, H)  # y_col = np.arange(0, height)
    y = np.tile(y_col, (W, 1)).T
    y = y.astype(np.float32)
    v_m_v0 = y - v0

    x = u_m_u0 / fx
    y = v_m_v0 / fy
    z = np.ones_like(x)
    pw = np.stack([x, y, z], axis=2)  # [h, w, c]
    return pw


def reconstruct_pcd(depth, fx, fy, u0, v0, pcd_base=None, mask=None):
    if type(depth) == torch.__name__:
        depth = depth.cpu().numpy().squeeze()
    depth = cv2.medianBlur(depth, 5)
    if pcd_base is None:
        H, W = depth.shape
        pcd_base = get_pcd_base(H, W, u0, v0, fx, fy)
    pcd = depth[:, :, None] * pcd_base
    if mask:
        pcd[mask] = 0
    return pcd


def save_point_cloud(pcd, rgb, filename, binary=True):
    """Save an RGB point cloud as a PLY file.
    :paras
        @pcd: Nx3 matrix, the XYZ coordinates
        @rgb: Nx3 matrix, the rgb colors for each 3D point
    """
    assert pcd.shape[0] == rgb.shape[0]

    if rgb is None:
        gray_concat = np.tile(np.array([128], dtype=np.uint8),
                              (pcd.shape[0], 3))
        points_3d = np.hstack((pcd, gray_concat))
    else:
        points_3d = np.hstack((pcd, rgb))
    python_types = (float, float, float, int, int, int)
    npy_types = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'),
                 ('green', 'u1'), ('blue', 'u1')]
    if binary is True:
        # Format into Numpy structured array
        vertices = []
        for row_idx in range(points_3d.shape[0]):
            cur_point = points_3d[row_idx]
            vertices.append(
                tuple(
                    dtype(point)
                    for dtype, point in zip(python_types, cur_point)))
        vertices_array = np.array(vertices, dtype=npy_types)
        el = PlyElement.describe(vertices_array, 'vertex')

        # write
        PlyData([el]).write(filename)
    else:
        x = np.squeeze(points_3d[:, 0])
        y = np.squeeze(points_3d[:, 1])
        z = np.squeeze(points_3d[:, 2])
        r = np.squeeze(points_3d[:, 3])
        g = np.squeeze(points_3d[:, 4])
        b = np.squeeze(points_3d[:, 5])

        ply_head = 'ply\n' \
                    'format ascii 1.0\n' \
                    'element vertex %d\n' \
                    'property float x\n' \
                    'property float y\n' \
                    'property float z\n' \
                    'property uchar red\n' \
                    'property uchar green\n' \
                    'property uchar blue\n' \
                    'end_header' % r.shape[0]
        # ---- Save ply data to disk
        np.savetxt(filename, np.column_stack[x, y, z, r, g, b], fmt='%f %f %f %d %d %d', header=ply_head, comments='')
mono/utils/unproj_pcd.py
hugoycj/Metric3D_MVP
import matplotlib.pyplot as plt
import os, cv2
import numpy as np
from mono.utils.transform import gray_to_colormap
import shutil
import glob
from mono.utils.running import main_process
import torch
from html4vision import Col, imagetable

def save_raw_imgs( 
    pred: torch.tensor,  
    rgb: torch.tensor, 
    filename: str, 
    save_dir: str,
    scale: float=200.0, 
    target: torch.tensor=None,
    ):
    """
    Save raw GT, predictions, RGB in the same file.
    """
    cv2.imwrite(os.path.join(save_dir, filename[:-4]+'_rgb.jpg'), rgb)
    cv2.imwrite(os.path.join(save_dir, filename[:-4]+'_d.png'), (pred*scale).astype(np.uint16))
    if target is not None:
        cv2.imwrite(os.path.join(save_dir, filename[:-4]+'_gt.png'), (target*scale).astype(np.uint16))
    

def save_val_imgs(
    iter: int, 
    pred: torch.tensor, 
    target: torch.tensor,
    rgb: torch.tensor, 
    filename: str, 
    save_dir: str, 
    tb_logger=None
    ):
    """
    Save GT, predictions, RGB in the same file.
    """
    rgb, pred_scale, target_scale, pred_color, target_color = get_data_for_log(pred, target, rgb)
    rgb = rgb.transpose((1, 2, 0))
    cat_img = np.concatenate([rgb, pred_color, target_color], axis=0)
    plt.imsave(os.path.join(save_dir, filename[:-4]+'_merge.jpg'), cat_img)

    # save to tensorboard
    if tb_logger is not None:
        tb_logger.add_image(f'{filename[:-4]}_merge.jpg', cat_img.transpose((2, 0, 1)), iter)

def get_data_for_log(pred: torch.tensor, target: torch.tensor, rgb: torch.tensor):
    mean = np.array([123.675, 116.28, 103.53])[:, np.newaxis, np.newaxis]
    std= np.array([58.395, 57.12, 57.375])[:, np.newaxis, np.newaxis]

    pred = pred.squeeze().cpu().numpy()
    target = target.squeeze().cpu().numpy()
    rgb = rgb.squeeze().cpu().numpy()

    pred[pred<0] = 0
    target[target<0] = 0
    max_scale = max(pred.max(), target.max())
    pred_scale = (pred/max_scale * 10000).astype(np.uint16)
    target_scale = (target/max_scale * 10000).astype(np.uint16)
    pred_color = gray_to_colormap(pred)
    target_color = gray_to_colormap(target)
    pred_color = cv2.resize(pred_color, (rgb.shape[2], rgb.shape[1]))
    target_color = cv2.resize(target_color, (rgb.shape[2], rgb.shape[1]))

    rgb = ((rgb * std) + mean).astype(np.uint8)
    return rgb, pred_scale, target_scale, pred_color, target_color


def create_html(name2path, save_path='index.html', size=(256, 384)):
    # table description
    cols = []
    for k, v in name2path.items():
        col_i =  Col('img', k, v) # specify image content for column
        cols.append(col_i)
    # html table generation
    imagetable(cols, out_file=save_path, imsize=size)


mono/utils/visualization.py
hugoycj/Metric3D_MVP