Models/__init__.py

"""
Description:
(1) 构建算法总体框架，包括输入双目图像的预处理、选用匹配网络、视差结果后处理
(2) input: stereo images
    output: disparity map & confidence (options)
(3) This file including some utils funcs
"""
import importlib
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from .tile import *


def lcm(x, y):
    xy_gcd = math.gcd(x, y)
    xy_lcm = x * y // xy_gcd
    return xy_lcm


__Model_Dict__ = {
    "UniFormer_FlowFormer": {
        "dr": lcm(16, 14),              # 最大降采样比例
        "pad_method": "resize",         # padding 方法
        "unpad_method": "restore",      # unpadding 方法
    },
    "stereobase": {
        "dr": 16,              # 最大降采样比例
        "pad_method": "resize",         # padding 方法
        "unpad_method": "restore",      # unpadding 方法
    },
}


# 导入模型, 格式: 文件夹 + 模型 (两者同名)
def import_model(model_name: str):
    # load the module, will raise ImportError if module cannot be loaded
    if model_name not in __Model_Dict__.keys():
        raise ValueError(
            f"Model {model_name} not in MODELS list. Valid models are {__Model_Dict__.keys()}"
        )
    m = importlib.import_module("Models." + model_name)     # 得到Models文件夹下的module文件
    return getattr(m, model_name)                           # 得到class的定义


class IMG_Processer(object):
    def __init__(self, d_rate):
        # padding方式
        self.rate = d_rate          # 最大降采样倍率

    def padding(self, img):
        # 图像padding, 左边和上边做padding
        _, _, H, W = img.shape
        h_pad = ((H // self.rate + 1) * self.rate - H) % self.rate
        w_pad = ((W // self.rate + 1) * self.rate - W) % self.rate
        ### Improving: top reflect for context preserve, right zero pad for imambugurous matching
        # self.size = (0, w_pad, h_pad, 0)
        # img = F.pad(img, pad=(0, w_pad, 0, 0), mode="constant")   # top right pad, usually for disparity
        # img = F.pad(img, pad=(0, 0, h_pad, 0), mode="reflect")
        
        ### following flow padding mode, e.g., used in FlowFormer, unimatch
        self.size = (w_pad//2, w_pad-w_pad//2, 0, h_pad)    # for kitti
        # self.size = (w_pad//2, w_pad-w_pad//2, h_pad//2, h_pad-h_pad//2)    # for sintel
        img = F.pad(img, pad=self.size, mode="replicate")
        return img

    def unpadding(self, img):
        # 去除padding区域
        h, w = img.shape[-2:]
        pad_wl, pad_wr, pad_ht, pad_hb = self.size
        img = img[..., pad_ht:h-pad_hb, pad_wl:w-pad_wr]
        return img

    def resize(self, img, update=True):
        # 将图像 resize 到合适的尺寸进行推理. 如果要对图像进行 resize, 需在 channel 前2维 padding 0
        _, _, H, W = img.shape
        if update:
            # 更新当前输入 batch 的原尺寸大小
            self.size = (H, W)
        resize_h = math.ceil(H / self.rate) * self.rate
        resize_w = math.ceil(W / self.rate) * self.rate
        img = F.interpolate(img, size=(resize_h, resize_w), mode="bilinear", align_corners=True)
        img[:, 0] *= resize_w / W 
        img[:, 1] *= resize_h / H
        return img

    def restore(self, flow):
        _, _, h, w = flow.shape
        flow = F.interpolate(flow, size=self.size, mode="bilinear", align_corners=True)
        (H, W) = self.size
        flow[:, 1] *= H / h
        flow[:, 0] *= W / w
        return flow


# 总模型定义
class MODEL(nn.Module):
    def __init__(self, config):
        super(MODEL, self).__init__()
        self.config = config
        model_name = config['model']['name']
        self.backbone = import_model(model_name)(config)
        self.model_name = model_name
        self.processer = IMG_Processer(d_rate=__Model_Dict__[model_name]["dr"])

    def forward(self, imgL, imgR, training_size=[384,384], flow_init=None):

        if self.training:
            output = self.backbone(imgL, imgR)

        else:
            pad_fn = getattr(self.processer, __Model_Dict__[self.model_name]["pad_method"])
            unpad_fn = getattr(self.processer, __Model_Dict__[self.model_name]["unpad_method"])
            if self.config['model']['tile_forward']:
                IMAGE_SIZE = imgL.shape[-2:]
                try:
                    TRAIN_SIZE = __Model_Dict__[self.model_name]["tile"]
                except:
                    TRAIN_SIZE = training_size
                # print("tile size is ", TRAIN_SIZE)
                TRAIN_SIZE[0] = min(TRAIN_SIZE[0], IMAGE_SIZE[0])  # 避免溢出
                TRAIN_SIZE[1] = min(TRAIN_SIZE[1], IMAGE_SIZE[1])

                min_overlap = (224,224)#(TRAIN_SIZE[0]-128, TRAIN_SIZE[1]-128)  # 重叠大小
                hws = compute_grid_indices(IMAGE_SIZE, TRAIN_SIZE, min_overlap=min_overlap)
                weights = compute_weight(hws, IMAGE_SIZE, TRAIN_SIZE, sigma=0.05)
                train_size = TRAIN_SIZE
                image_size = IMAGE_SIZE

                print(f"Using tiling strategy, tile size: {train_size}, overlap: {min_overlap}")

                # 基于全图(全局)预测先验, 再进行 tile 细化估计 (效果提升不明显, 先验作为初值并不能被网络很好地感知和利用起来)
                flows = 0
                flow_count = 0
                for idx, (h, w) in enumerate(hws):
                    image1_tile = imgL[:, :, h:h + train_size[0], w:w + train_size[1]]  # 截取图像
                    image2_tile = imgR[:, :, h:h + train_size[0], w:w + train_size[1]]

                    # 如果存在光流先验
                    if flow_init is not None:
                        flow_tile = flow_init[:, :, h:h + train_size[0], w:w + train_size[1]]
                        all_tile_pad = pad_fn(torch.cat([flow_tile, image1_tile, image2_tile], dim=1))
                        flow_tile_pad, image1_tile_pad, image2_tile_pad = all_tile_pad[:,0:2], all_tile_pad[:,2:5], all_tile_pad[:,5:8]
                    else:
                        flow_tile = torch.zeros_like(image1_tile)[:,:2]
                        all_tile_pad = pad_fn(torch.cat([flow_tile, image1_tile, image2_tile], dim=1))
                        flow_tile_pad, image1_tile_pad, image2_tile_pad = None, all_tile_pad[:,2:5], all_tile_pad[:,5:8]
                    
                    output = self.backbone(image1_tile_pad, image2_tile_pad, flow_tile_pad)            # 对 padding 后的截取图像进行推理
                    flow_pred_pad = output['flow']
                    if "conf" in output.keys():
                        flow_conf_pad = output['conf']
                        flow_comb_pad = torch.cat([flow_pred_pad, flow_conf_pad], dim=1)
                        flow_comb = unpad_fn(flow_comb_pad)                 # 去 padding
                        flow_pre  = flow_comb[:, :2]
                        flow_conf = flow_comb[:, 2:]
                        # 置信度设置最小阈值, 避免后续除零导致 NAN
                        flow_conf = flow_conf.clamp(min=1e-6)
                    else:
                        flow_pre = unpad_fn(flow_pred_pad)                  # 去 padding
                        flow_conf = weights[idx]

                    padding = (w, image_size[1] - w - train_size[1], h, image_size[0] - h - train_size[0], 0, 0)
                    flows += F.pad(flow_pre * flow_conf, padding)
                    flow_count += F.pad(flow_conf, padding)                          # 计算当前区域的光流和加权权重

                flow_pre = flows / flow_count

                # output = {"flow": flow_pre, "conf": flow_count}
                output = {"flow": flow_pre, "conf": torch.ones_like(flow_pre)[:,:1]}

            else:
                # 输入图像预处理, padding以适合网络
                imgL = pad_fn(imgL)
                imgR = pad_fn(imgR)
                output = self.backbone(imgL, imgR)

                # 对 padding 区域进行后处理
                # if self.config['task'] == "stereo":
                #     disp = output['disparity']
                #     zero = torch.zeros_like(disp)
                #     flow = torch.cat([disp, zero], dim=1)
                #     flow = self.processer.restore(flow)
                #     output['disparity'] = - flow[:, 0:1]
                #
                # elif self.config['task'] == "flow":
                output['flow'] = unpad_fn(output['flow'])
                # output['conf'] = F.interpolate(output['conf'], size=(output['flow'].shape[-2:]), mode="bilinear", align_corners=True) if 'conf' in output.keys() else torch.ones_like(output['flow'])[:,:1]
                output['conf'] = torch.ones_like(output['flow'])[:,:1]

            # else:
            #     raise NotImplementedError
            # TODO: 增加视差后处理模块(基于交叉熵的置信度估计滤波)
        return output


class LOSS(nn.Module):
    def __init__(self, config):
        super(LOSS, self).__init__()
        self.config = config
        model_name = config['model']['name']
        # 动态导入对应模型的损失函数模块
        m = importlib.import_module("Models." + model_name)
        loss_module = getattr(m, "loss_func")
        self.loss = loss_module(config)

    def forward(self, label_dict, training_output):
        return self.loss(label_dict, training_output)