UPFlow_pytorch/model/upflow.py

from __future__ import absolute_import, division, print_function
import collections
import torch
import torch.nn as nn
import torch.nn.functional as F
# from torch.nn.utils.spectral_norm import spectral_norm
from model.pwc_modules import conv, initialize_msra, upsample2d_flow_as, upsample_flow, FlowEstimatorDense_v2, ContextNetwork_v2_, OccEstimatorDense, OccContextNetwork
from model.pwc_modules import WarpingLayer_no_div, FeatureExtractor
from model.correlation_package.correlation import Correlation
import numpy as np
from utils.tools import tools
from utils.loss import loss_functions
from utils.pytorch_correlation import Corr_pyTorch
import cv2
import os
import math


class network_tools():
    class sgu_model(tools.abstract_model):
        def __init__(self):
            super(network_tools.sgu_model, self).__init__()

            class FlowEstimatorDense_temp(tools.abstract_model):

                def __init__(self, ch_in, f_channels=(128, 128, 96, 64, 32), ch_out=2):
                    super(FlowEstimatorDense_temp, self).__init__()
                    N = 0
                    ind = 0
                    N += ch_in
                    self.conv1 = conv(N, f_channels[ind])
                    N += f_channels[ind]

                    ind += 1
                    self.conv2 = conv(N, f_channels[ind])
                    N += f_channels[ind]

                    ind += 1
                    self.conv3 = conv(N, f_channels[ind])
                    N += f_channels[ind]

                    ind += 1
                    self.conv4 = conv(N, f_channels[ind])
                    N += f_channels[ind]

                    ind += 1
                    self.conv5 = conv(N, f_channels[ind])
                    N += f_channels[ind]
                    self.num_feature_channel = N
                    ind += 1
                    self.conv_last = conv(N, ch_out, isReLU=False)

                def forward(self, x):
                    x1 = torch.cat([self.conv1(x), x], dim=1)
                    x2 = torch.cat([self.conv2(x1), x1], dim=1)
                    x3 = torch.cat([self.conv3(x2), x2], dim=1)
                    x4 = torch.cat([self.conv4(x3), x3], dim=1)
                    x5 = torch.cat([self.conv5(x4), x4], dim=1)
                    x_out = self.conv_last(x5)
                    return x5, x_out

            f_channels_es = (32, 32, 32, 16, 8)
            in_C = 64
            self.warping_layer = WarpingLayer_no_div()
            self.dense_estimator_mask = FlowEstimatorDense_temp(in_C, f_channels=f_channels_es, ch_out=3)
            self.upsample_output_conv = nn.Sequential(conv(3, 16, kernel_size=3, stride=1, dilation=1),
                                                      conv(16, 16, stride=2),
                                                      conv(16, 32, kernel_size=3, stride=1, dilation=1),
                                                      conv(32, 32, stride=2), )

        def forward(self, flow_init, feature_1, feature_2, output_level_flow=None):
            n, c, h, w = flow_init.shape
            n_f, c_f, h_f, w_f = feature_1.shape
            if h != h_f or w != w_f:
                flow_init = upsample2d_flow_as(flow_init, feature_1, mode="bilinear", if_rate=True)
            feature_2_warp = self.warping_layer(feature_2, flow_init)
            input_feature = torch.cat((feature_1, feature_2_warp), dim=1)
            feature, x_out = self.dense_estimator_mask(input_feature)
            inter_flow = x_out[:, :2, :, :]
            inter_mask = x_out[:, 2, :, :]
            inter_mask = torch.unsqueeze(inter_mask, 1)
            inter_mask = torch.sigmoid(inter_mask)
            n_, c_, h_, w_ = inter_flow.shape
            if output_level_flow is not None:
                inter_flow = upsample2d_flow_as(inter_flow, output_level_flow, mode="bilinear", if_rate=True)
                inter_mask = upsample2d_flow_as(inter_mask, output_level_flow, mode="bilinear")
                flow_init = output_level_flow
            flow_up = tools.torch_warp(flow_init, inter_flow) * (1 - inter_mask) + flow_init * inter_mask
            return flow_init, flow_up, inter_flow, inter_mask

        def output_conv(self, x):
            return self.upsample_output_conv(x)

    @classmethod
    def normalize_features(cls, feature_list, normalize, center, moments_across_channels=True, moments_across_images=True):
        """Normalizes feature tensors (e.g., before computing the cost volume).
        Args:
          feature_list: list of torch tensors, each with dimensions [b, c, h, w]
          normalize: bool flag, divide features by their standard deviation
          center: bool flag, subtract feature mean
          moments_across_channels: bool flag, compute mean and std across channels, 看到UFlow默认是True
          moments_across_images: bool flag, compute mean and std across images, 看到UFlow默认是True

        Returns:
          list, normalized feature_list
        """

        # Compute feature statistics.

        statistics = collections.defaultdict(list)
        axes = [1, 2, 3] if moments_across_channels else [2, 3]  # [b, c, h, w]
        for feature_image in feature_list:
            mean = torch.mean(feature_image, dim=axes, keepdim=True)  # [b,1,1,1] or [b,c,1,1]
            variance = torch.var(feature_image, dim=axes, keepdim=True)  # [b,1,1,1] or [b,c,1,1]
            statistics['mean'].append(mean)
            statistics['var'].append(variance)

        if moments_across_images:
            # statistics['mean'] = ([tf.reduce_mean(input_tensor=statistics['mean'])] *
            #                       len(feature_list))
            # statistics['var'] = [tf.reduce_mean(input_tensor=statistics['var'])
            #                      ] * len(feature_list)
            statistics['mean'] = ([torch.mean(torch.stack(statistics['mean'], dim=0), dim=(0,))] * len(feature_list))
            statistics['var'] = ([torch.var(torch.stack(statistics['var'], dim=0), dim=(0,))] * len(feature_list))

        statistics['std'] = [torch.sqrt(v + 1e-16) for v in statistics['var']]

        # Center and normalize features.

        if center:
            feature_list = [
                f - mean for f, mean in zip(feature_list, statistics['mean'])
            ]
        if normalize:
            feature_list = [f / std for f, std in zip(feature_list, statistics['std'])]

        return feature_list

    @classmethod
    def weighted_ssim(cls, x, y, weight, c1=float('inf'), c2=9e-6, weight_epsilon=0.01):
        """Computes a weighted structured image similarity measure.
        Args:
          x: a batch of images, of shape [B, C, H, W].
          y:  a batch of images, of shape [B, C, H, W].
          weight: shape [B, 1, H, W], representing the weight of each
            pixel in both images when we come to calculate moments (means and
            correlations). values are in [0,1]
          c1: A floating point number, regularizes division by zero of the means.
          c2: A floating point number, regularizes division by zero of the second
            moments.
          weight_epsilon: A floating point number, used to regularize division by the
            weight.

        Returns:
          A tuple of two pytorch Tensors. First, of shape [B, C, H-2, W-2], is scalar
          similarity loss per pixel per channel, and the second, of shape
          [B, 1, H-2. W-2], is the average pooled `weight`. It is needed so that we
          know how much to weigh each pixel in the first tensor. For example, if
          `'weight` was very small in some area of the images, the first tensor will
          still assign a loss to these pixels, but we shouldn't take the result too
          seriously.
        """

        def _avg_pool3x3(x):
            # tf kernel [b,h,w,c]
            return F.avg_pool2d(x, (3, 3), (1, 1))
            # return tf.nn.avg_pool(x, [1, 3, 3, 1], [1, 1, 1, 1], 'VALID')

        if c1 == float('inf') and c2 == float('inf'):
            raise ValueError('Both c1 and c2 are infinite, SSIM loss is zero. This is '
                             'likely unintended.')
        average_pooled_weight = _avg_pool3x3(weight)
        weight_plus_epsilon = weight + weight_epsilon
        inverse_average_pooled_weight = 1.0 / (average_pooled_weight + weight_epsilon)

        def weighted_avg_pool3x3(z):
            wighted_avg = _avg_pool3x3(z * weight_plus_epsilon)
            return wighted_avg * inverse_average_pooled_weight

        mu_x = weighted_avg_pool3x3(x)
        mu_y = weighted_avg_pool3x3(y)
        sigma_x = weighted_avg_pool3x3(x ** 2) - mu_x ** 2
        sigma_y = weighted_avg_pool3x3(y ** 2) - mu_y ** 2
        sigma_xy = weighted_avg_pool3x3(x * y) - mu_x * mu_y
        if c1 == float('inf'):
            ssim_n = (2 * sigma_xy + c2)
            ssim_d = (sigma_x + sigma_y + c2)
        elif c2 == float('inf'):
            ssim_n = 2 * mu_x * mu_y + c1
            ssim_d = mu_x ** 2 + mu_y ** 2 + c1
        else:
            ssim_n = (2 * mu_x * mu_y + c1) * (2 * sigma_xy + c2)
            ssim_d = (mu_x ** 2 + mu_y ** 2 + c1) * (sigma_x + sigma_y + c2)
        result = ssim_n / ssim_d
        return torch.clamp((1 - result) / 2, 0, 1), average_pooled_weight

    @classmethod
    def edge_aware_smoothness_order1(cls, img, pred):
        def gradient_x(img):
            gx = img[:, :, :-1, :] - img[:, :, 1:, :]
            return gx

        def gradient_y(img):
            gy = img[:, :, :, :-1] - img[:, :, :, 1:]
            return gy

        pred_gradients_x = gradient_x(pred)
        pred_gradients_y = gradient_y(pred)

        image_gradients_x = gradient_x(img)
        image_gradients_y = gradient_y(img)

        weights_x = torch.exp(-torch.mean(torch.abs(image_gradients_x), 1, keepdim=True))
        weights_y = torch.exp(-torch.mean(torch.abs(image_gradients_y), 1, keepdim=True))

        smoothness_x = torch.abs(pred_gradients_x) * weights_x
        smoothness_y = torch.abs(pred_gradients_y) * weights_y
        return torch.mean(smoothness_x) + torch.mean(smoothness_y)

    @classmethod
    def edge_aware_smoothness_order2(cls, img, pred):
        def gradient_x(img, stride=1):
            gx = img[:, :, :-stride, :] - img[:, :, stride:, :]
            return gx

        def gradient_y(img, stride=1):
            gy = img[:, :, :, :-stride] - img[:, :, :, stride:]
            return gy

        pred_gradients_x = gradient_x(pred)
        pred_gradients_xx = gradient_x(pred_gradients_x)
        pred_gradients_y = gradient_y(pred)
        pred_gradients_yy = gradient_y(pred_gradients_y)

        image_gradients_x = gradient_x(img, stride=2)
        image_gradients_y = gradient_y(img, stride=2)

        weights_x = torch.exp(-torch.mean(torch.abs(image_gradients_x), 1, keepdim=True))
        weights_y = torch.exp(-torch.mean(torch.abs(image_gradients_y), 1, keepdim=True))

        smoothness_x = torch.abs(pred_gradients_xx) * weights_x
        smoothness_y = torch.abs(pred_gradients_yy) * weights_y
        return torch.mean(smoothness_x) + torch.mean(smoothness_y)

    @classmethod
    def flow_smooth_delta(cls, flow, if_second_order=False):
        def gradient(x):
            D_dy = x[:, :, 1:] - x[:, :, :-1]
            D_dx = x[:, :, :, 1:] - x[:, :, :, :-1]
            return D_dx, D_dy

        dx, dy = gradient(flow)
        # dx2, dxdy = gradient(dx)
        # dydx, dy2 = gradient(dy)
        if if_second_order:
            dx2, dxdy = gradient(dx)
            dydx, dy2 = gradient(dy)
            smooth_loss = dx.abs().mean() + dy.abs().mean() + dx2.abs().mean() + dxdy.abs().mean() + dydx.abs().mean() + dy2.abs().mean()
        else:
            smooth_loss = dx.abs().mean() + dy.abs().mean()
        # smooth_loss = dx.abs().mean() + dy.abs().mean()  # + dx2.abs().mean() + dxdy.abs().mean() + dydx.abs().mean() + dy2.abs().mean()
        # 暂时不上二阶的平滑损失，似乎加上以后就太猛了，无法降低photo loss TODO
        return smooth_loss

    @classmethod
    def photo_loss_multi_type(cls, x, y, occ_mask, photo_loss_type='abs_robust',  # abs_robust, charbonnier,L1, SSIM
                              photo_loss_delta=0.4, photo_loss_use_occ=False,
                              ):
        occ_weight = occ_mask
        if photo_loss_type == 'abs_robust':
            photo_diff = x - y
            loss_diff = (torch.abs(photo_diff) + 0.01).pow(photo_loss_delta)
        elif photo_loss_type == 'charbonnier':
            photo_diff = x - y
            loss_diff = ((photo_diff) ** 2 + 1e-6).pow(photo_loss_delta)
        elif photo_loss_type == 'L1':
            photo_diff = x - y
            loss_diff = torch.abs(photo_diff + 1e-6)
        elif photo_loss_type == 'SSIM':
            loss_diff, occ_weight = cls.weighted_ssim(x, y, occ_mask)
        else:
            raise ValueError('wrong photo_loss type: %s' % photo_loss_type)

        if photo_loss_use_occ:
            photo_loss = torch.sum(loss_diff * occ_weight) / (torch.sum(occ_weight) + 1e-6)
        else:
            photo_loss = torch.mean(loss_diff)
        return photo_loss


class UPFlow_net(tools.abstract_model):
    class config(tools.abstract_config):
        def __init__(self):
            # occ loss choose
            self.occ_type = 'for_back_check'
            self.alpha_1 = 0.1
            self.alpha_2 = 0.5
            self.occ_check_obj_out_all = 'obj'  # if boundary dilated warping is used, here should be obj
            self.stop_occ_gradient = False
            self.smooth_level = 'final'  # final or 1/4
            self.smooth_type = 'edge'  # edge or delta
            self.smooth_order_1_weight = 1
            # smooth loss
            self.smooth_order_2_weight = 0
            # photo loss type add SSIM
            self.photo_loss_type = 'abs_robust'  # abs_robust, charbonnier,L1, SSIM
            self.photo_loss_delta = 0.4
            self.photo_loss_use_occ = False
            self.photo_loss_census_weight = 0
            # use cost volume norm
            self.if_norm_before_cost_volume = False
            self.norm_moments_across_channels = True
            self.norm_moments_across_images = True
            self.multi_scale_distillation_weight = 0
            self.multi_scale_distillation_style = 'upup'  # down,upup,
            # 'down', 'upup', 'updown'
            self.multi_scale_distillation_occ = True  # if consider occlusion mask in multiscale distilation
            self.if_froze_pwc = False
            self.input_or_sp_input = 1  # use raw input or special input for photo loss
            self.if_use_boundary_warp = True  # if use the boundary dilated warping

            self.if_sgu_upsample = False  # if use sgu upsampling
            self.if_use_cor_pytorch = False  # use my implementation of correlation layer by pytorch. only for test model in cpu(corr layer cuda is not compiled)

        def __call__(self, ):
            # return PWCNet_unsup_irr_bi_v5_4(self)
            return UPFlow_net(self)

    def __init__(self, conf: config):
        super(UPFlow_net, self).__init__()
        # === get config file
        self.conf = conf

        # === build the network
        self.search_range = 4
        self.num_chs = [3, 16, 32, 64, 96, 128, 196]
        #                  1/2 1/4 1/8 1/16 1/32 1/64
        self.estimator_f_channels = (128, 128, 96, 64, 32)
        self.context_f_channels = (128, 128, 128, 96, 64, 32, 2)
        self.output_level = 4
        self.num_levels = 7
        self.leakyRELU = nn.LeakyReLU(0.1, inplace=True)
        self.feature_pyramid_extractor = FeatureExtractor(self.num_chs)
        self.warping_layer = WarpingLayer_no_div()
        self.dim_corr = (self.search_range * 2 + 1) ** 2
        self.num_ch_in = self.dim_corr + 32 + 2
        self.flow_estimators = FlowEstimatorDense_v2(self.num_ch_in, f_channels=self.estimator_f_channels)
        self.context_networks = ContextNetwork_v2_(self.flow_estimators.n_channels + 2, f_channels=self.context_f_channels)
        self.conv_1x1 = nn.ModuleList([conv(196, 32, kernel_size=1, stride=1, dilation=1),
                                       conv(128, 32, kernel_size=1, stride=1, dilation=1),
                                       conv(96, 32, kernel_size=1, stride=1, dilation=1),
                                       conv(64, 32, kernel_size=1, stride=1, dilation=1),
                                       conv(32, 32, kernel_size=1, stride=1, dilation=1)])
        self.occ_check_model_ls = []
        self.correlation_pytorch = Corr_pyTorch(pad_size=self.search_range, kernel_size=1,
                                                max_displacement=self.search_range, stride1=1, stride2=1)  # correlation layer using pytorch
        # === build sgu upsampling
        if self.conf.if_sgu_upsample:
            self.sgi_model = network_tools.sgu_model()
        else:
            self.sgi_model = None

        # === build loss function
        self.occ_check_model = tools.occ_check_model(occ_type=self.conf.occ_type, occ_alpha_1=self.conf.alpha_1, occ_alpha_2=self.conf.alpha_2,
                                                     obj_out_all=self.conf.occ_check_obj_out_all)
        initialize_msra(self.modules())
        if self.conf.if_froze_pwc:
            self.froze_PWC()

    def forward(self, input_dict: dict):
        '''
        :param input_dict:     im1, im2, im1_raw, im2_raw, start, if_loss
        :return: output_dict:  flows, flow_f_out, flow_b_out, photo_loss
        '''
        im1_ori, im2_ori = input_dict['im1'], input_dict['im2']  # in training: the cropped image; in testing: the input image
        if input_dict['if_loss']:
            if self.conf.input_or_sp_input == 1:
                im1, im2 = im1_ori, im2_ori
            else:
                im1, im2 = input_dict['im1_sp'], input_dict['im2_sp']  # change the input image to special input image and the original images are used for loss computing
        else:
            im1, im2 = im1_ori, im2_ori

        output_dict = {}
        flow_f_pwc_out, flow_b_pwc_out, flows = self.forward_2_frame_v3(im1, im2, if_loss=input_dict['if_loss'])  # forward estimation
        occ_fw, occ_bw = self.occ_check_model(flow_f=flow_f_pwc_out, flow_b=flow_b_pwc_out)  # 0 in occ area, 1 in others

        '''  ======================================    =====================================  '''
        output_dict['flow_f_out'] = flow_f_pwc_out
        output_dict['flow_b_out'] = flow_b_pwc_out
        output_dict['occ_fw'] = occ_fw
        output_dict['occ_bw'] = occ_bw

        if input_dict['if_loss']:
            # === smooth loss
            if self.conf.smooth_level == 'final':
                s_flow_f, s_flow_b = flow_f_pwc_out, flow_b_pwc_out
                s_im1, s_im2 = im1_ori, im2_ori
            elif self.conf.smooth_level == '1/4':
                s_flow_f, s_flow_b = flows[0]  # flow in 1/4 scale
                _, _, temp_h, temp_w = s_flow_f.size()
                s_im1 = F.interpolate(im1_ori, (temp_h, temp_w), mode='area')
                s_im2 = F.interpolate(im2_ori, (temp_h, temp_w), mode='area')
            else:
                raise ValueError('wrong smooth level choosed: %s' % self.smooth_level)
            smooth_loss = 0
            # 1 order smooth loss
            if self.conf.smooth_order_1_weight > 0:
                if self.conf.smooth_type == 'edge':
                    smooth_loss += self.conf.smooth_order_1_weight * network_tools.edge_aware_smoothness_order1(img=s_im1, pred=s_flow_f)
                    smooth_loss += self.conf.smooth_order_1_weight * network_tools.edge_aware_smoothness_order1(img=s_im2, pred=s_flow_b)
                elif self.conf.smooth_type == 'delta':
                    smooth_loss += self.conf.smooth_order_1_weight * network_tools.flow_smooth_delta(flow=s_flow_f, if_second_order=False)
                    smooth_loss += self.conf.smooth_order_1_weight * network_tools.flow_smooth_delta(flow=s_flow_b, if_second_order=False)
                else:
                    raise ValueError('wrong smooth_type: %s' % self.conf.smooth_type)

            # 2 order smooth loss
            if self.conf.smooth_order_2_weight > 0:
                if self.conf.smooth_type == 'edge':
                    smooth_loss += self.conf.smooth_order_2_weight * network_tools.edge_aware_smoothness_order2(img=s_im1, pred=s_flow_f)
                    smooth_loss += self.conf.smooth_order_2_weight * network_tools.edge_aware_smoothness_order2(img=s_im2, pred=s_flow_b)
                elif self.conf.smooth_type == 'delta':
                    smooth_loss += self.conf.smooth_order_2_weight * network_tools.flow_smooth_delta(flow=s_flow_f, if_second_order=True)
                    smooth_loss += self.conf.smooth_order_2_weight * network_tools.flow_smooth_delta(flow=s_flow_b, if_second_order=True)
                else:
                    raise ValueError('wrong smooth_type: %s' % self.conf.smooth_type)
            output_dict['smooth_loss'] = smooth_loss

            # === photo loss
            if self.conf.if_use_boundary_warp:
                im1_s, im2_s, start_s = input_dict['im1_raw'], input_dict['im2_raw'], input_dict['start']  # the image before cropping
                im1_warp = tools.boundary_dilated_warp.warp_im(im2_s, flow_f_pwc_out, start_s)  # warped im1 by forward flow and im2
                im2_warp = tools.boundary_dilated_warp.warp_im(im1_s, flow_b_pwc_out, start_s)
            else:
                im1_warp = tools.torch_warp(im2_ori, flow_f_pwc_out)  # warped im1 by forward flow and im2
                im2_warp = tools.torch_warp(im1_ori, flow_b_pwc_out)
            # photo loss
            if self.conf.stop_occ_gradient:
                occ_fw, occ_bw = occ_fw.clone().detach(), occ_bw.clone().detach()
            photo_loss = network_tools.photo_loss_multi_type(im1_ori, im1_warp, occ_fw, photo_loss_type=self.conf.photo_loss_type,
                                                             photo_loss_delta=self.conf.photo_loss_delta, photo_loss_use_occ=self.conf.photo_loss_use_occ)
            photo_loss += network_tools.photo_loss_multi_type(im2_ori, im2_warp, occ_bw, photo_loss_type=self.conf.photo_loss_type,
                                                              photo_loss_delta=self.conf.photo_loss_delta, photo_loss_use_occ=self.conf.photo_loss_use_occ)
            output_dict['photo_loss'] = photo_loss
            output_dict['im1_warp'] = im1_warp
            output_dict['im2_warp'] = im2_warp

            # === census loss
            if self.conf.photo_loss_census_weight > 0:
                census_loss = loss_functions.census_loss_torch(img1=im1_ori, img1_warp=im1_warp, mask=occ_fw, q=self.conf.photo_loss_delta,
                                                               charbonnier_or_abs_robust=False, if_use_occ=self.conf.photo_loss_use_occ, averge=True) + \
                              loss_functions.census_loss_torch(img1=im2_ori, img1_warp=im2_warp, mask=occ_bw, q=self.conf.photo_loss_delta,
                                                               charbonnier_or_abs_robust=False, if_use_occ=self.conf.photo_loss_use_occ, averge=True)
                census_loss *= self.conf.photo_loss_census_weight
            else:
                census_loss = None
            output_dict['census_loss'] = census_loss

            # === multi scale distillation loss
            if self.conf.multi_scale_distillation_weight > 0:
                flow_fw_label = flow_f_pwc_out.clone().detach()
                flow_bw_label = flow_b_pwc_out.clone().detach()
                msd_loss_ls = []
                for i, (scale_fw, scale_bw) in enumerate(flows):
                    if self.conf.multi_scale_distillation_style == 'down':
                        flow_fw_label_sacle = upsample_flow(flow_fw_label, target_flow=scale_fw)
                        occ_scale_fw = F.interpolate(occ_fw, [scale_fw.size(2), scale_fw.size(3)], mode='nearest')
                        flow_bw_label_sacle = upsample_flow(flow_bw_label, target_flow=scale_bw)
                        occ_scale_bw = F.interpolate(occ_bw, [scale_bw.size(2), scale_bw.size(3)], mode='nearest')
                    elif self.conf.multi_scale_distillation_style == 'upup':
                        flow_fw_label_sacle = flow_fw_label
                        scale_fw = upsample_flow(scale_fw, target_flow=flow_fw_label_sacle)
                        occ_scale_fw = occ_fw
                        flow_bw_label_sacle = flow_bw_label
                        scale_bw = upsample_flow(scale_bw, target_flow=flow_bw_label_sacle)
                        occ_scale_bw = occ_bw
                    else:
                        raise ValueError('wrong multi_scale_distillation_style: %s' % self.conf.multi_scale_distillation_style)
                    msd_loss_scale_fw = network_tools.photo_loss_multi_type(x=scale_fw, y=flow_fw_label_sacle, occ_mask=occ_scale_fw, photo_loss_type='abs_robust',
                                                                            photo_loss_use_occ=self.conf.multi_scale_distillation_occ)
                    msd_loss_ls.append(msd_loss_scale_fw)
                    msd_loss_scale_bw = network_tools.photo_loss_multi_type(x=scale_bw, y=flow_bw_label_sacle, occ_mask=occ_scale_bw, photo_loss_type='abs_robust',
                                                                            photo_loss_use_occ=self.conf.multi_scale_distillation_occ)
                    msd_loss_ls.append(msd_loss_scale_bw)
                msd_loss = sum(msd_loss_ls)
                msd_loss = self.conf.multi_scale_distillation_weight * msd_loss
            else:
                msd_loss = None

            output_dict['msd_loss'] = msd_loss
        return output_dict

    def forward_2_frame_v3(self, x1_raw, x2_raw, if_loss=False):
        _, _, height_im, width_im = x1_raw.size()
        # on the bottom level are original images
        x1_pyramid = self.feature_pyramid_extractor(x1_raw) + [x1_raw]
        x2_pyramid = self.feature_pyramid_extractor(x2_raw) + [x2_raw]
        flows = []
        # init
        b_size, _, h_x1, w_x1, = x1_pyramid[0].size()
        init_dtype = x1_pyramid[0].dtype
        init_device = x1_pyramid[0].device
        flow_f = torch.zeros(b_size, 2, h_x1, w_x1, dtype=init_dtype, device=init_device).float()
        flow_b = torch.zeros(b_size, 2, h_x1, w_x1, dtype=init_dtype, device=init_device).float()
        # build pyramid
        feature_level_ls = []
        for l, (x1, x2) in enumerate(zip(x1_pyramid, x2_pyramid)):
            x1_1by1 = self.conv_1x1[l](x1)
            x2_1by1 = self.conv_1x1[l](x2)
            feature_level_ls.append((x1, x1_1by1, x2, x2_1by1))
            if l == self.output_level:
                break
        for level, (x1, x1_1by1, x2, x2_1by1) in enumerate(feature_level_ls):
            flow_f, flow_b, flow_f_res, flow_b_res = self.decode_level_res(level=level, flow_1=flow_f, flow_2=flow_b,
                                                                           feature_1=x1, feature_1_1x1=x1_1by1,
                                                                           feature_2=x2, feature_2_1x1=x2_1by1,
                                                                           img_ori_1=x1_raw, img_ori_2=x2_raw)
            flow_f = flow_f + flow_f_res
            flow_b = flow_b + flow_b_res
            flows.append([flow_f, flow_b])
        flow_f_out = upsample2d_flow_as(flow_f, x1_raw, mode="bilinear", if_rate=True)
        flow_b_out = upsample2d_flow_as(flow_b, x1_raw, mode="bilinear", if_rate=True)

        # === upsample to full resolution
        if self.conf.if_sgu_upsample:
            feature_1_1x1 = self.sgi_model.output_conv(x1_raw)
            feature_2_1x1 = self.sgi_model.output_conv(x2_raw)
            flow_f_out = self.self_guided_upsample(flow_up_bilinear=flow_f, feature_1=feature_1_1x1, feature_2=feature_2_1x1, output_level_flow=flow_f_out)
            flow_b_out = self.self_guided_upsample(flow_up_bilinear=flow_b, feature_1=feature_2_1x1, feature_2=feature_1_1x1, output_level_flow=flow_b_out)
        else:
            pass
        return flow_f_out, flow_b_out, flows[::-1]

    def decode_level_res(self, level, flow_1, flow_2, feature_1, feature_1_1x1, feature_2, feature_2_1x1, img_ori_1, img_ori_2):
        flow_1_up_bilinear = upsample2d_flow_as(flow_1, feature_1, mode="bilinear", if_rate=True)
        flow_2_up_bilinear = upsample2d_flow_as(flow_2, feature_2, mode="bilinear", if_rate=True)
        # warping
        if level == 0:
            feature_2_warp = feature_2
            feature_1_warp = feature_1
        else:
            if self.conf.if_sgu_upsample:
                flow_1_up_bilinear = self.self_guided_upsample(flow_up_bilinear=flow_1_up_bilinear, feature_1=feature_1_1x1, feature_2=feature_2_1x1)
                flow_2_up_bilinear = self.self_guided_upsample(flow_up_bilinear=flow_2_up_bilinear, feature_1=feature_2_1x1, feature_2=feature_1_1x1)
            feature_2_warp = self.warping_layer(feature_2, flow_1_up_bilinear)
            feature_1_warp = self.warping_layer(feature_1, flow_2_up_bilinear)
        # if norm feature
        if self.conf.if_norm_before_cost_volume:
            feature_1, feature_2_warp = network_tools.normalize_features((feature_1, feature_2_warp), normalize=True, center=True,
                                                                         moments_across_channels=self.conf.norm_moments_across_channels,
                                                                         moments_across_images=self.conf.norm_moments_across_images)
            feature_2, feature_1_warp = network_tools.normalize_features((feature_2, feature_1_warp), normalize=True, center=True,
                                                                         moments_across_channels=self.conf.norm_moments_across_channels,
                                                                         moments_across_images=self.conf.norm_moments_across_images)
        # correlation
        if self.conf.if_use_cor_pytorch:
            out_corr_1 = self.correlation_pytorch(feature_1, feature_2_warp)
            out_corr_2 = self.correlation_pytorch(feature_2, feature_1_warp)
        else:
            out_corr_1 = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(feature_1, feature_2_warp)
            out_corr_2 = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(feature_2, feature_1_warp)
        out_corr_relu_1 = self.leakyRELU(out_corr_1)
        out_corr_relu_2 = self.leakyRELU(out_corr_2)
        feature_int_1, flow_res_1 = self.flow_estimators(torch.cat([out_corr_relu_1, feature_1_1x1, flow_1_up_bilinear], dim=1))
        feature_int_2, flow_res_2 = self.flow_estimators(torch.cat([out_corr_relu_2, feature_2_1x1, flow_2_up_bilinear], dim=1))
        flow_1_up_bilinear_ = flow_1_up_bilinear + flow_res_1
        flow_2_up_bilinear_ = flow_2_up_bilinear + flow_res_2
        flow_fine_1 = self.context_networks(torch.cat([feature_int_1, flow_1_up_bilinear_], dim=1))
        flow_fine_2 = self.context_networks(torch.cat([feature_int_2, flow_2_up_bilinear_], dim=1))
        flow_1_res = flow_res_1 + flow_fine_1
        flow_2_res = flow_res_2 + flow_fine_2
        return flow_1_up_bilinear, flow_2_up_bilinear, flow_1_res, flow_2_res

    def froze_PWC(self):
        for param in self.feature_pyramid_extractor.parameters():
            param.requires_grad = False
        for param in self.flow_estimators.parameters():
            param.requires_grad = False
        for param in self.context_networks.parameters():
            param.requires_grad = False
        for param in self.conv_1x1.parameters():
            param.requires_grad = False

    def self_guided_upsample(self, flow_up_bilinear, feature_1, feature_2, output_level_flow=None):
        flow_up_bilinear_, out_flow, inter_flow, inter_mask = self.sgi_model(flow_up_bilinear, feature_1, feature_2, output_level_flow=output_level_flow)
        return out_flow

    @classmethod
    def demo(cls):
        param_dict = {
            'occ_type': 'for_back_check',
            'alpha_1': 0.1,
            'alpha_2': 0.5,
            'occ_check_obj_out_all': 'obj',
            'stop_occ_gradient': False,
            'smooth_level': 'final',  # final or 1/4
            'smooth_type': 'edge',  # edge or delta
            'smooth_order_1_weight': 1,
            # smooth loss
            'smooth_order_2_weight': 0,
            # photo loss type add SSIM
            'photo_loss_type': 'abs_robust',  # abs_robust, charbonnier,L1, SSIM
            'photo_loss_delta': 0.4,
            'photo_loss_use_occ': False,
            'photo_loss_census_weight': 1,
            # use cost volume norm
            'if_norm_before_cost_volume': True,
            'norm_moments_across_channels': False,
            'norm_moments_across_images': False,
            'multi_scale_distillation_weight': 1,
            'multi_scale_distillation_style': 'upup',
            'multi_scale_photo_weight': 1,  # 'down', 'upup', 'updown'
            'multi_scale_distillation_occ': True,  # if consider occlusion mask in multiscale distilation
            'if_froze_pwc': False,
            'input_or_sp_input': 1,
            'if_use_boundary_warp': True,
            'if_use_cor_pytorch': True,
        }
        net_conf = UPFlow_net.config()
        net_conf.update(param_dict)
        net_conf.get_name(print_now=True)
        net = net_conf()  # .cuda()
        net.eval()
        im = np.random.random((1, 3, 320, 320))
        start = np.zeros((1, 2, 1, 1))
        start = torch.from_numpy(start).float()  # .cuda()
        im_torch = torch.from_numpy(im).float()  # .cuda()
        input_dict = {'im1': im_torch, 'im2': im_torch,
                      'im1_raw': im_torch, 'im2_raw': im_torch, 'im1_sp': im_torch, 'im2_sp': im_torch, 'start': start, 'if_loss': True}
        output_dict = net(input_dict)
        print('smooth_loss', output_dict['smooth_loss'], 'photo_loss', output_dict['photo_loss'], 'census_loss', output_dict['census_loss'])
        for i in output_dict.keys():
            if output_dict[i] is None:
                print(i, output_dict[i])
            else:
                tools.check_tensor(output_dict[i], i)


if __name__ == '__main__':
    UPFlow_net.demo()