warping.py

# Taken from https://github.com/modelscope/modelscope

# The implementation here is modified based on flow-style-vton,
# originally Apache 2.0 License and publicly available at https://github.com/SenHe/Flow-Style-VTON

from math import sqrt

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F



def apply_offset(offset):
    sizes = list(offset.size()[2:])
    grid_list = torch.meshgrid(
        [torch.arange(size, device=offset.device) for size in sizes])
    grid_list = reversed(grid_list)
    grid_list = [
        grid.float().unsqueeze(0) + offset[:, dim, ...]
        for dim, grid in enumerate(grid_list)
    ]
    grid_list = [
        grid / ((size - 1.0) / 2.0) - 1.0
        for grid, size in zip(grid_list, reversed(sizes))
    ]

    return torch.stack(grid_list, dim=-1)


class Conv2dBlock(nn.Module):

    def __init__(self,
                 input_dim,
                 output_dim,
                 kernel_size,
                 stride,
                 padding=0,
                 norm='batch',
                 activation='prelu',
                 pad_type='zero',
                 bias=True):
        super().__init__()
        self.use_bias = bias
        # initialize padding
        if pad_type == 'reflect':
            self.pad = nn.ReflectionPad2d(padding)
        elif pad_type == 'replicate':
            self.pad = nn.ReplicationPad2d(padding)
        elif pad_type == 'zero':
            self.pad = nn.ZeroPad2d(padding)
        else:
            assert 0, 'Unsupported padding type: {}'.format(pad_type)

        norm_dim = input_dim
        if norm == 'batch':
            self.norm = nn.BatchNorm2d(norm_dim)
        elif norm == 'instance':
            self.norm = nn.InstanceNorm2d(norm_dim)
        elif norm == 'none':
            self.norm = None
        else:
            assert 0, 'Unsupported normalization: {}'.format(norm)

        # initialize activation
        if activation == 'relu':
            self.activation = nn.ReLU(inplace=True)
        elif activation == 'lrelu':
            self.activation = nn.LeakyReLU(0.1, inplace=True)
        elif activation == 'prelu':
            self.activation = nn.PReLU()
        elif activation == 'selu':
            self.activation = nn.SELU(inplace=True)
        elif activation == 'tanh':
            self.activation = nn.Tanh()
        elif activation == 'sigmoid':
            self.activation = nn.Sigmoid()

        elif activation == 'none':
            self.activation = None
        else:
            assert 0, 'Unsupported activation: {}'.format(activation)

        # initialize convolution
        self.conv = nn.Conv2d(
            input_dim, output_dim, kernel_size, stride, bias=self.use_bias)

    def forward(self, x):
        if self.norm:
            x = self.norm(x)
        if self.activation:
            x = self.activation(x)
        x = self.conv(self.pad(x))
        return x


class ResBlock(nn.Module):

    def __init__(self, in_channels):
        super(ResBlock, self).__init__()

        self.block = nn.Sequential(
            nn.BatchNorm2d(in_channels), nn.PReLU(),
            nn.Conv2d(
                in_channels, in_channels, kernel_size=3, padding=1,
                bias=False), nn.BatchNorm2d(in_channels), nn.PReLU(),
            nn.Conv2d(
                in_channels, in_channels, kernel_size=3, padding=1,
                bias=False))

    def forward(self, x):
        return self.block(x) + x


class ResBlock_2(nn.Module):

    def __init__(self, in_channels):
        super(ResBlock_2, self).__init__()
        self.block = nn.Sequential(
            nn.BatchNorm2d(in_channels), nn.PReLU(),
            nn.Conv2d(
                in_channels, in_channels, kernel_size=3, padding=1,
                bias=False), nn.BatchNorm2d(in_channels), nn.PReLU(),
            nn.Conv2d(
                in_channels, in_channels, kernel_size=3, padding=1,
                bias=False))

    def forward(self, x1, x2):
        return self.block(x1) + x2


class DownSample(nn.Module):

    def __init__(self, in_channels, out_channels):
        super(DownSample, self).__init__()
        self.block = nn.Sequential(
            nn.BatchNorm2d(in_channels), nn.PReLU(),
            nn.Conv2d(
                in_channels,
                out_channels,
                kernel_size=3,
                stride=2,
                padding=1,
                bias=False))

    def forward(self, x):
        return self.block(x)


class UpSample(nn.Module):

    def __init__(self, in_channels, out_channels):
        super(UpSample, self).__init__()
        self.block = nn.Sequential(
            nn.BatchNorm2d(in_channels), nn.PReLU(),
            nn.ConvTranspose2d(
                in_channels, out_channels, kernel_size=2, stride=2,
                bias=False))

    def forward(self, x):
        return self.block(x)

    def upsample(self, F, scale):
        """[2x nearest neighbor upsampling]
        Arguments:
            F {[torch tensor]} -- [tensor to be upsampled, (B, C, H, W)]
        """
        upsample = torch.nn.Upsample(scale_factor=scale, mode='nearest')
        return upsample(F)


class FeatureEncoder(nn.Module):

    def __init__(self, in_channels, chns=[64, 128, 256, 256, 256]):
        # in_channels = 3 for images, and is larger (e.g., 17+1+1) for agnositc representation
        super(FeatureEncoder, self).__init__()
        self.encoders = []
        for i, out_chns in enumerate(chns):
            if i == 0:
                encoder = nn.Sequential(
                    DownSample(in_channels, out_chns), ResBlock(out_chns),
                    ResBlock(out_chns))
            else:
                encoder = nn.Sequential(
                    DownSample(chns[i - 1], out_chns), ResBlock(out_chns),
                    ResBlock(out_chns))

            self.encoders.append(encoder)

        self.encoders = nn.ModuleList(self.encoders)

    def forward(self, x):
        encoder_features = []
        for encoder in self.encoders:
            x = encoder(x)
            encoder_features.append(x)
        return encoder_features


class RefinePyramid(nn.Module):

    def __init__(self, chns=[64, 128, 256, 256, 256], fpn_dim=256):
        super(RefinePyramid, self).__init__()
        self.chns = chns

        # adaptive
        self.adaptive = []
        for in_chns in list(reversed(chns)):
            adaptive_layer = nn.Conv2d(in_chns, fpn_dim, kernel_size=1)
            self.adaptive.append(adaptive_layer)
        self.adaptive = nn.ModuleList(self.adaptive)
        # output conv
        self.smooth = []
        for i in range(len(chns)):
            smooth_layer = nn.Conv2d(
                fpn_dim, fpn_dim, kernel_size=3, padding=1)
            self.smooth.append(smooth_layer)
        self.smooth = nn.ModuleList(self.smooth)

    def forward(self, x):
        conv_ftr_list = x

        feature_list = []
        last_feature = None
        for i, conv_ftr in enumerate(list(reversed(conv_ftr_list))):
            # adaptive
            feature = self.adaptive[i](conv_ftr)
            # fuse
            if last_feature is not None:
                feature = feature + F.interpolate(
                    last_feature, scale_factor=2, mode='nearest')
            # smooth
            feature = self.smooth[i](feature)
            last_feature = feature
            feature_list.append(feature)

        return tuple(reversed(feature_list))


def Round(x):
    ''' x: tensor '''
    return torch.round(x) - x.detach() + x


def Camp(x, mi=0, ma=192):
    ''' x: tensor '''
    if x < mi:
        x = (mi - x.detach()) + x
    elif x > ma:
        x = x - (x.detach() - ma)
    return x


class CorrelationLayer(nn.Module):

    def __init__(self, init_scale=5):
        super(CorrelationLayer, self).__init__()
        self.init_scale = init_scale
        self.softmax3 = nn.Softmax(dim=3)
        self.sig = nn.Sigmoid()
        self.reduce1 = Conv2dBlock(
            64, 4, 1, 1, 0, activation='prelu', norm='none')
        self.reduce2 = Conv2dBlock(
            64, 4, 1, 1, 0, activation='prelu', norm='none')

        init_h, init_w = 192 / self.init_scale, 256 / self.init_scale
        self.step_h, self.step_w = [], []
        step_h_half = [i + 1 for i in range(0, int(init_h + 1), 2)]
        step_h_half_re = list(reversed(step_h_half[1:]))
        self.step_h.extend(step_h_half_re)
        self.step_h.extend(step_h_half)
        step_w_half = [i + 1 for i in range(0, int(init_w + 1), 2)]
        step_w_half_re = list(reversed(step_w_half[1:]))
        self.step_w.extend(step_w_half_re)
        self.step_w.extend(step_w_half)
        self.mask_h, self.mask_w = len(self.step_h), len(self.step_w)

    def forward(self, location, fea_c, fea_p, scale_param, H, W, c_landmark,
                p_landmark):
        init_h, init_w = torch.tensor(
            W / self.init_scale, device=fea_c.device), torch.tensor(
                H / self.init_scale, device=fea_c.device)
        N, C, fea_H = fea_c.shape[0], fea_c.shape[1], fea_c.shape[2]
        downsample_ratio = H / fea_H
        landmark_flow = -1 * torch.ones((N, 64, H, W), device=fea_c.device)
        mask_h, mask_w = self.mask_h, self.mask_w

        fea_cn = self.upsample(fea_c, scale=downsample_ratio)
        fea_pn = self.upsample(fea_p, scale=downsample_ratio)

        fea_cn = self.reduce1(fea_cn)
        fea_pn = self.reduce2(fea_pn)
        C = fea_cn.shape[1]

        src_box_h0 = torch.tensor(
            self.step_h,
            device=fea_cn.device).unsqueeze(dim=0).unsqueeze(dim=0).repeat(
                N, 32, 1)
        src_box_w0 = torch.tensor(
            self.step_w,
            device=fea_cn.device).unsqueeze(dim=0).unsqueeze(dim=0).repeat(
                N, 32, 1)
        tar_box_h0 = torch.tensor(
            self.step_h,
            device=fea_cn.device).unsqueeze(dim=0).unsqueeze(dim=0).repeat(
                N, 32, 1)
        tar_box_w0 = torch.tensor(
            self.step_w,
            device=fea_cn.device).unsqueeze(dim=0).unsqueeze(dim=0).repeat(
                N, 32, 1)
        adj_cw, adj_ch, adj_pw, adj_ph = scale_param[:, :, 0, 0], scale_param[:, :, 0, 1],\
            scale_param[:, :, 1, 0], scale_param[:, :, 1, 1]
        c_h, c_w = (init_h * self.sig(adj_ch)).unsqueeze(
            dim=2), (init_w * self.sig(adj_cw)).unsqueeze(dim=2)
        p_h, p_w = (init_h * self.sig(adj_ph)).unsqueeze(
            dim=2), (init_w * self.sig(adj_pw)).unsqueeze(dim=2)
        src_box_h = self.sig((c_h - src_box_h0) * 2)
        src_box_w = self.sig((c_w - src_box_w0) * 2)
        tar_box_h = self.sig((p_h - tar_box_h0) * 2)
        tar_box_w = self.sig((p_w - tar_box_w0) * 2)

        src_box_h, src_box_w = src_box_h.unsqueeze(dim=3), src_box_w.unsqueeze(
            dim=2)
        tar_box_h, tar_box_w = tar_box_h.unsqueeze(dim=3), tar_box_w.unsqueeze(
            dim=2)
        src_mask = torch.matmul(src_box_h, src_box_w)
        tar_mask = torch.matmul(tar_box_h, tar_box_w)

        cloth_patch_all = torch.zeros((N, 32 * C, mask_h, mask_w),
                                      device=fea_cn.device)
        person_patch_all = torch.zeros((N, 32 * C, mask_h, mask_w),
                                       device=fea_cn.device)

        location_patch_all = -1 * torch.ones(
            (N, 64, mask_h, mask_w), device=fea_cn.device)
        one_flow_all = -1 * torch.ones(
            (N, 64, mask_h, mask_w), device=fea_cn.device)
        coord_info = []
        mask = torch.zeros((N, 32, H, W), device=fea_cn.device)

        for b in range(N):
            for i in range(32):
                c_center_x, c_center_y = c_landmark[b, i, 0], c_landmark[b, i,
                                                                         1]
                p_center_x, p_center_y = p_landmark[b, i, 0], p_landmark[b, i,
                                                                         1]
                if (c_center_x == 0
                        and c_center_y == 0) or (p_center_x == 0
                                                 and p_center_y == 0):
                    p_new_top_y, p_new_bottom_y, p_new_left_x, p_new_right_x, p_new_h, p_new_w = 0, 0, 0, 0, 0, 0
                else:
                    c_left_x = torch.floor(c_center_x.int() - mask_w / 2)
                    c_right_x = torch.floor(c_center_x.int() + mask_w / 2)
                    delta_x1 = int(torch.clamp(0 - c_left_x, min=0))
                    delta_x2 = int(torch.clamp(c_right_x - W, min=0))
                    c_left_x, c_right_x = int(c_left_x), int(c_right_x)
                    c_top_y = torch.floor(c_center_y.int() - mask_h / 2)
                    c_bottom_y = torch.floor(c_center_y.int() + mask_h / 2)
                    delta_y1 = int(torch.clamp(0 - c_top_y, min=0))
                    delta_y2 = int(torch.clamp(c_bottom_y - H, min=0))
                    c_top_y, c_bottom_y = int(c_top_y), int(c_bottom_y)

                    p_left_x = torch.floor(p_center_x.int() - mask_w / 2)
                    p_right_x = torch.floor(p_center_x.int() + mask_w / 2)
                    delta_x3 = int(torch.clamp(0 - p_left_x, min=0))
                    delta_x4 = int(torch.clamp(p_right_x - W, min=0))
                    p_left_x, p_right_x = int(p_left_x), int(p_right_x)
                    p_top_y = torch.floor(p_center_y.int() - mask_h / 2)
                    p_bottom_y = torch.floor(p_center_y.int() + mask_h / 2)
                    delta_y3 = int(torch.clamp(0 - p_top_y, min=0))
                    delta_y4 = int(torch.clamp(p_bottom_y - H, min=0))
                    p_top_y, p_bottom_y = int(p_top_y), int(p_bottom_y)

                    # Deformable Patch
                    c_new_top_y, c_new_bottom_y, c_new_left_x, c_new_right_x = \
                        c_top_y + delta_y1, c_bottom_y - delta_y2, c_left_x + delta_x1, c_right_x - delta_x2
                    p_new_top_y, p_new_bottom_y, p_new_left_x, p_new_right_x = \
                        p_top_y + delta_y3, p_bottom_y - delta_y4, p_left_x + delta_x3, p_right_x - delta_x4
                    c_new_h, c_new_w = c_new_bottom_y - c_new_top_y, c_new_right_x - c_new_left_x
                    p_new_h, p_new_w = p_new_bottom_y - p_new_top_y, p_new_right_x - p_new_left_x

                    cloth_patch_all[b, i * C:(i + 1) * C, : c_new_h, : c_new_w] = \
                        fea_cn[b, :, c_new_top_y: c_new_bottom_y, c_new_left_x: c_new_right_x] * \
                        src_mask[b, i, 0 + delta_y1: mask_h - delta_y2, 0 + delta_x1: mask_w - delta_x2]
                    person_patch_all[b, i * C: (i + 1) * C, : p_new_h, : p_new_w] = \
                        fea_pn[b, :, p_new_top_y: p_new_bottom_y, p_new_left_x: p_new_right_x] * \
                        tar_mask[b, i, 0 + delta_y3: mask_h - delta_y4, 0 + delta_x3: mask_w - delta_x4]
                    location_patch_all[b, i * 2:i * 2
                                       + 2, :c_new_h, :c_new_w] = location[
                                           b, i * 2:i * 2 + 2,
                                           c_new_top_y:c_new_bottom_y,
                                           c_new_left_x:c_new_right_x]

                coord_info.append([
                    p_new_top_y, p_new_bottom_y, p_new_left_x, p_new_right_x,
                    p_new_h, p_new_w
                ])

        Q = person_patch_all.view(N, 32, C, mask_h,
                                  mask_w).view(N, 32, C,
                                               -1).permute(0, 1, 3,
                                                           2).contiguous()
        K = cloth_patch_all.view(N, 32, C, mask_h, mask_w).view(N, 32, C, -1)
        correlation = self.softmax3(torch.matmul(Q, K))
        V = location_patch_all.view(N, 32, 2, mask_h,
                                    mask_w).view(N, 32, 2,
                                                 -1).permute(0, 1, 3,
                                                             2).contiguous()
        one_flow_all = torch.matmul(correlation,
                                    V).permute(0, 1, 3, 2).contiguous().view(
                                        N, 32 * 2, mask_h, mask_w)

        for b in range(N):
            for i in range(32):
                p_new_top_y, p_new_bottom_y, p_new_left_x, p_new_right_x, p_new_h, p_new_w = coord_info[
                    b * 32 + i]
                if p_new_h != 0 and p_new_w != 0:
                    landmark_flow[b, i * 2: i * 2 + 2, p_new_top_y: p_new_bottom_y, p_new_left_x: p_new_right_x] = \
                        one_flow_all[b, i * 2: i * 2 + 2, : p_new_h, : p_new_w]
                    mask[b, i, p_new_top_y:p_new_bottom_y,
                         p_new_left_x:p_new_right_x] = torch.tensor(
                             1, device=fea_cn.device)

        return landmark_flow, mask

    def upsample(self, F, scale):
        """[2x nearest neighbor upsampling]
        Arguments:
            F {[torch tensor]} -- [tensor to be upsampled, (B, C, H, W)]
        """
        upsample = torch.nn.Upsample(scale_factor=scale, mode='nearest')
        return upsample(F)


class LocalFlow(nn.Module):

    def __init__(self, fpn_dim=256, use_num=3, init_scale=6):
        super(LocalFlow, self).__init__()
        self.use_num = use_num
        self.fusers = []
        self.fusers_refine = []
        self.location_preds = []
        self.patch_preds = []
        self.map_preds = []
        self.map_refine = []
        self.reduc = []

        for i in range(use_num):
            in_chns = fpn_dim
            map_refine = nn.Conv2d(
                32, 32, kernel_size=3, stride=1, padding=1, groups=32)

            reduc = nn.Sequential(
                Conv2dBlock(
                    in_chns, 64, 3, 1, 1, activation='prelu', bias=False))
            fusers = nn.Sequential(
                Conv2dBlock(
                    fpn_dim * 2,
                    fpn_dim,
                    3,
                    1,
                    1,
                    activation='prelu',
                    bias=False),
                Conv2dBlock(
                    fpn_dim, 32, 3, 1, 1, activation='prelu', bias=False))
            fusers_refine_layer = [
                nn.Upsample(
                    scale_factor=2, mode='bilinear', align_corners=True),
                ResBlock(32)
            ] * (
                i + 1)
            fusers_refine = nn.Sequential(*fusers_refine_layer)
            location_preds = nn.Sequential(
                Conv2dBlock(32, 64, 3, 1, 1, activation='prelu', norm='none'),
                Conv2dBlock(64, 64, 3, 1, 1, activation='prelu', norm='none'),
                Conv2dBlock(64, 64, 3, 1, 1, activation='prelu', norm='none'),
            )
            # patch
            patch_layer = [DownSample(32, 32), ResBlock(32)] * (use_num + 2)
            patch_layer.append(nn.AdaptiveAvgPool2d((2, 2)))
            patch_preds = nn.Sequential(*patch_layer)
            # attention map
            map_layer = []
            map_layer = [
                ResBlock(32),
                Conv2dBlock(32, 32, 3, 1, 1, activation='prelu', norm='none')
            ]
            map_preds = nn.Sequential(*map_layer)

            self.reduc.append(reduc)
            self.fusers.append(fusers)
            self.fusers_refine.append(fusers_refine)
            self.location_preds.append(location_preds)
            self.patch_preds.append(patch_preds)
            self.map_preds.append(map_preds)
            self.map_refine.append(map_refine)

        self.reduc = nn.ModuleList(self.reduc)

        self.fusers = nn.ModuleList(self.fusers)
        self.fusers_refine = nn.ModuleList(self.fusers_refine)
        self.location_preds = nn.ModuleList(self.location_preds)
        self.patch_preds = nn.ModuleList(self.patch_preds)
        self.map_preds = nn.ModuleList(self.map_preds)
        self.map_refine = nn.ModuleList(self.map_refine)
        self.CorrelationLayer = CorrelationLayer(init_scale=init_scale)

    def forward(self, cloth, fea_c, fea_p, c_landmark, p_landmark):
        flow_list, map_list = [], []
        N, _, H, W = cloth.shape
        for i in range(self.use_num):
            fuse = self.fusers[i](torch.cat((fea_c[i], fea_p[i]), dim=1))
            fuse = self.fusers_refine[i](fuse)
            location = self.location_preds[i](fuse)
            patch = self.patch_preds[i](fuse)
            att_map = self.map_preds[i](fuse)
            fea_c_reduc = self.reduc[i](fea_c[i])
            fea_p_reduc = self.reduc[i](fea_p[i])
            flow, mask = self.CorrelationLayer(location, fea_c_reduc,
                                               fea_p_reduc, patch, H, W,
                                               c_landmark, p_landmark)
            if i != 0:
                flow = (flow + last_flow) / 2
                add_map = torch.add(last_att_map, att_map)
                att_map = self.map_refine[i](add_map)

            last_flow = flow
            flow_list.append(last_flow)
            last_att_map = att_map * mask
            map_list.append(last_att_map)

        return flow_list, map_list


class EqualLR:

    def __init__(self, name):
        self.name = name

    def compute_weight(self, module):
        weight = getattr(module, self.name + '_orig')
        fan_in = weight.data.size(1) * weight.data[0][0].numel()

        return weight * sqrt(2 / fan_in)

    @staticmethod
    def apply(module, name):
        fn = EqualLR(name)

        weight = getattr(module, name)
        del module._parameters[name]
        module.register_parameter(name + '_orig', nn.Parameter(weight.data))
        module.register_forward_pre_hook(fn)

        return fn

    def __call__(self, module, input):
        weight = self.compute_weight(module)
        setattr(module, self.name, weight)


def equal_lr(module, name='weight'):
    EqualLR.apply(module, name)

    return module


class EqualLinear(nn.Module):

    def __init__(self, in_dim, out_dim):
        super().__init__()

        linear = nn.Linear(in_dim, out_dim)
        linear.weight.data.normal_()
        linear.bias.data.zero_()

        self.linear = equal_lr(linear)

    def forward(self, input):
        return self.linear(input)


class ModulatedConv2d(nn.Module):

    def __init__(self,
                 fin,
                 fout,
                 kernel_size,
                 padding_type='zero',
                 upsample=False,
                 downsample=False,
                 latent_dim=512,
                 normalize_mlp=False):
        super(ModulatedConv2d, self).__init__()
        self.in_channels = fin
        self.out_channels = fout
        self.kernel_size = kernel_size
        padding_size = kernel_size // 2

        if kernel_size == 1:
            self.demudulate = False
        else:
            self.demudulate = True

        self.weight = nn.Parameter(
            torch.Tensor(fout, fin, kernel_size, kernel_size))
        self.bias = nn.Parameter(torch.Tensor(1, fout, 1, 1))

        if normalize_mlp:
            self.mlp_class_std = nn.Sequential(
                EqualLinear(latent_dim, fin), PixelNorm())
        else:
            self.mlp_class_std = EqualLinear(latent_dim, fin)

        if padding_type == 'reflect':
            self.padding = nn.ReflectionPad2d(padding_size)
        else:
            self.padding = nn.ZeroPad2d(padding_size)

        self.weight.data.normal_()
        self.bias.data.zero_()

    def forward(self, input, latent):
        fan_in = self.weight.data.size(1) * self.weight.data[0][0].numel()
        weight = self.weight * sqrt(2 / fan_in)
        weight = weight.view(1, self.out_channels, self.in_channels,
                             self.kernel_size, self.kernel_size)

        s = self.mlp_class_std(latent).view(-1, 1, self.in_channels, 1, 1)
        weight = s * weight
        if self.demudulate:
            d = torch.rsqrt((weight**2).sum(4).sum(3).sum(2) + 1e-5).view(
                -1, self.out_channels, 1, 1, 1)
            weight = (d * weight).view(-1, self.in_channels, self.kernel_size,
                                       self.kernel_size)
        else:
            weight = weight.view(-1, self.in_channels, self.kernel_size,
                                 self.kernel_size)

        batch, _, height, width = input.shape

        input = input.view(1, -1, height, width)
        input = self.padding(input)
        out = F.conv2d(
            input, weight, groups=batch).view(batch, self.out_channels, height,
                                              width) + self.bias

        return out


class StyledConvBlock(nn.Module):

    def __init__(self,
                 fin,
                 fout,
                 latent_dim=256,
                 padding='zero',
                 actvn='lrelu',
                 normalize_affine_output=False,
                 modulated_conv=False):
        super(StyledConvBlock, self).__init__()
        if not modulated_conv:
            if padding == 'reflect':
                padding_layer = nn.ReflectionPad2d
            else:
                padding_layer = nn.ZeroPad2d

        if modulated_conv:
            conv2d = ModulatedConv2d
        else:
            conv2d = EqualConv2d

        if modulated_conv:
            self.actvn_gain = sqrt(2)
        else:
            self.actvn_gain = 1.0

        self.modulated_conv = modulated_conv

        if actvn == 'relu':
            activation = nn.ReLU(True)
        else:
            activation = nn.LeakyReLU(0.2, True)

        if self.modulated_conv:
            self.conv0 = conv2d(
                fin,
                fout,
                kernel_size=3,
                padding_type=padding,
                upsample=False,
                latent_dim=latent_dim,
                normalize_mlp=normalize_affine_output)
        else:
            conv0 = conv2d(fin, fout, kernel_size=3)

            seq0 = [padding_layer(1), conv0]
            self.conv0 = nn.Sequential(*seq0)

        self.actvn0 = activation

        if self.modulated_conv:
            self.conv1 = conv2d(
                fout,
                fout,
                kernel_size=3,
                padding_type=padding,
                downsample=False,
                latent_dim=latent_dim,
                normalize_mlp=normalize_affine_output)
        else:
            conv1 = conv2d(fout, fout, kernel_size=3)
            seq1 = [padding_layer(1), conv1]
            self.conv1 = nn.Sequential(*seq1)

        self.actvn1 = activation

    def forward(self, input, latent=None):
        if self.modulated_conv:
            out = self.conv0(input, latent)
        else:
            out = self.conv0(input)

        out = self.actvn0(out) * self.actvn_gain

        if self.modulated_conv:
            out = self.conv1(out, latent)
        else:
            out = self.conv1(out)

        out = self.actvn1(out) * self.actvn_gain

        return out


class Styled_F_ConvBlock(nn.Module):

    def __init__(self,
                 fin,
                 fout,
                 latent_dim=256,
                 padding='zero',
                 actvn='prelu',
                 normalize_affine_output=False,
                 modulated_conv=False):
        super(Styled_F_ConvBlock, self).__init__()
        if not modulated_conv:
            if padding == 'reflect':
                padding_layer = nn.ReflectionPad2d
            else:
                padding_layer = nn.ZeroPad2d

        if modulated_conv:
            conv2d = ModulatedConv2d
        else:
            conv2d = EqualConv2d

        if modulated_conv:
            self.actvn_gain = sqrt(2)
        else:
            self.actvn_gain = 1.0

        self.modulated_conv = modulated_conv

        if actvn == 'relu':
            activation = nn.ReLU(True)
        else:
            activation = nn.LeakyReLU(0.2, True)

        if self.modulated_conv:
            self.conv0 = conv2d(
                fin,
                128,
                kernel_size=3,
                padding_type=padding,
                upsample=False,
                latent_dim=latent_dim,
                normalize_mlp=normalize_affine_output)
        else:
            conv0 = conv2d(fin, 128, kernel_size=3)

            seq0 = [padding_layer(1), conv0]
            self.conv0 = nn.Sequential(*seq0)

        self.actvn0 = activation

        if self.modulated_conv:
            self.conv1 = conv2d(
                128,
                fout,
                kernel_size=3,
                padding_type=padding,
                downsample=False,
                latent_dim=latent_dim,
                normalize_mlp=normalize_affine_output)
        else:
            conv1 = conv2d(128, fout, kernel_size=3)
            seq1 = [padding_layer(1), conv1]
            self.conv1 = nn.Sequential(*seq1)

    def forward(self, input, latent=None):
        if self.modulated_conv:
            out = self.conv0(input, latent)
        else:
            out = self.conv0(input)

        out = self.actvn0(out) * self.actvn_gain

        if self.modulated_conv:
            out = self.conv1(out, latent)
        else:
            out = self.conv1(out)

        return out


class AFlowNet(nn.Module):

    def __init__(self, num_pyramid, fpn_dim=256, use_num=3):
        super(AFlowNet, self).__init__()

        padding_type = 'zero'
        actvn = 'lrelu'
        normalize_mlp = False
        modulated_conv = True
        self.use_num = use_num

        self.netRefine = []
        self.netStyle = []
        self.netF = []
        self.map_preds = []
        self.map_refine = []

        for i in range(num_pyramid):

            netRefine_layer = torch.nn.Sequential(
                torch.nn.Conv2d(
                    2 * fpn_dim,
                    out_channels=128,
                    kernel_size=3,
                    stride=1,
                    padding=1), nn.PReLU(), ResBlock(128), ResBlock(128),
                torch.nn.Conv2d(
                    in_channels=128,
                    out_channels=64,
                    kernel_size=3,
                    stride=1,
                    padding=1), nn.PReLU(),
                torch.nn.Conv2d(
                    in_channels=64,
                    out_channels=32,
                    kernel_size=3,
                    stride=1,
                    padding=1), nn.PReLU(),
                torch.nn.Conv2d(
                    in_channels=32,
                    out_channels=12,
                    kernel_size=3,
                    stride=1,
                    padding=1), nn.PReLU(),
                torch.nn.Conv2d(
                    in_channels=12,
                    out_channels=2,
                    kernel_size=3,
                    stride=1,
                    padding=1))

            style_block = torch.nn.Sequential(
                torch.nn.Conv2d(
                    2 * fpn_dim,
                    out_channels=fpn_dim,
                    kernel_size=3,
                    stride=1,
                    padding=1), nn.PReLU(),
                torch.nn.Conv2d(
                    fpn_dim,
                    out_channels=128,
                    kernel_size=3,
                    stride=1,
                    padding=1), nn.PReLU(),
                torch.nn.Conv2d(
                    in_channels=128,
                    out_channels=64,
                    kernel_size=3,
                    stride=1,
                    padding=1), nn.PReLU(),
                torch.nn.Conv2d(
                    in_channels=64,
                    out_channels=32,
                    kernel_size=3,
                    stride=1,
                    padding=1), nn.PReLU(),
                torch.nn.Conv2d(
                    in_channels=32,
                    out_channels=2,
                    kernel_size=3,
                    stride=1,
                    padding=1))

            style_F_block = Styled_F_ConvBlock(
                49,
                2,
                latent_dim=256,
                padding=padding_type,
                actvn=actvn,
                normalize_affine_output=normalize_mlp,
                modulated_conv=modulated_conv)

            map_preds = torch.nn.Sequential(
                torch.nn.Conv2d(
                    2 * fpn_dim,
                    out_channels=128,
                    kernel_size=3,
                    stride=1,
                    padding=1), nn.PReLU(),
                torch.nn.Conv2d(
                    in_channels=128,
                    out_channels=64,
                    kernel_size=3,
                    stride=1,
                    padding=1), nn.PReLU(),
                torch.nn.Conv2d(
                    in_channels=64,
                    out_channels=32,
                    kernel_size=3,
                    stride=1,
                    padding=1), nn.PReLU(),
                torch.nn.Conv2d(
                    in_channels=32,
                    out_channels=2,
                    kernel_size=3,
                    stride=1,
                    padding=1))
            if i == 0:
                map_refine = nn.Sequential(
                    Conv2dBlock(
                        2, 2, 3, 1, 1, activation='prelu', norm='none'),
                    Conv2dBlock(
                        2, 2, 3, 1, 1, activation='prelu', norm='none'))
            elif i >= 2:
                map_refine = nn.Sequential(
                    Conv2dBlock(
                        36, 16, 3, 1, 1, activation='prelu', norm='none'),
                    Conv2dBlock(
                        16, 4, 3, 1, 1, activation='prelu', norm='none'),
                    Conv2dBlock(
                        4, 2, 3, 1, 1, activation='prelu', norm='none'))
            else:
                map_refine = nn.Sequential(
                    Conv2dBlock(
                        4, 2, 3, 1, 1, activation='prelu', norm='none'),
                    Conv2dBlock(
                        2, 2, 3, 1, 1, activation='prelu', norm='none'))

            self.netRefine.append(netRefine_layer)
            self.netStyle.append(style_block)
            self.netF.append(style_F_block)
            self.map_preds.append(map_preds)
            self.map_refine.append(map_refine)

        self.netRefine = nn.ModuleList(self.netRefine)
        self.netStyle = nn.ModuleList(self.netStyle)
        self.netF = nn.ModuleList(self.netF)
        self.map_preds = nn.ModuleList(self.map_preds)
        self.map_refine = nn.ModuleList(self.map_refine)

        self.cond_style = torch.nn.Sequential(
            torch.nn.Conv2d(256, 128, kernel_size=(8, 6), stride=1, padding=0),
            nn.PReLU())
        self.image_style = torch.nn.Sequential(
            torch.nn.Conv2d(256, 128, kernel_size=(8, 6), stride=1, padding=0),
            nn.PReLU())

    def forward(self, global_flow_input, device, warp_feature=True):

        x = global_flow_input[0]
        x_edge = global_flow_input[1]
        x_warps = global_flow_input[2]
        x_conds = global_flow_input[3]
        localmap_list = global_flow_input[5]
        last_flow = None
        last_gmap = None
        gmap_all = []
        last_flow_all = []
        delta_list = []
        x_all = []
        x_edge_all = []
        cond_fea_all = []
        delta_x_all = []
        delta_y_all = []
        filter_x = [[0, 0, 0], [1, -2, 1], [0, 0, 0]]
        filter_y = [[0, 1, 0], [0, -2, 0], [0, 1, 0]]
        filter_diag1 = [[1, 0, 0], [0, -2, 0], [0, 0, 1]]
        filter_diag2 = [[0, 0, 1], [0, -2, 0], [1, 0, 0]]
        weight_array = np.ones([3, 3, 1, 4])
        weight_array[:, :, 0, 0] = filter_x
        weight_array[:, :, 0, 1] = filter_y
        weight_array[:, :, 0, 2] = filter_diag1
        weight_array[:, :, 0, 3] = filter_diag2

        weight_array = torch.FloatTensor(weight_array).permute(3, 2, 0,
                                                               1).to(device)
        self.weight = nn.Parameter(data=weight_array, requires_grad=False)

        fea_num = len(x_warps)

        for i in range(fea_num):
            x_warp = x_warps[fea_num - 1 - i]
            x_cond = x_conds[fea_num - 1 - i]
            cond_fea_all.append(x_cond)

            if last_flow is not None and warp_feature:
                x_warp_after = F.grid_sample(
                    x_warp,
                    last_flow.detach().permute(0, 2, 3, 1),
                    mode='bilinear',
                    padding_mode='border')
            else:
                x_warp_after = x_warp

            flow = self.netStyle[i](torch.cat([x_warp_after, x_cond], 1))
            delta_list.append(flow)
            flow = apply_offset(flow)
            if last_flow is not None:
                flow = F.grid_sample(
                    last_flow, flow, mode='bilinear', padding_mode='border')
            else:
                flow = flow.permute(0, 3, 1, 2)

            last_flow = flow
            x_warp = F.grid_sample(
                x_warp,
                flow.permute(0, 2, 3, 1),
                mode='bilinear',
                padding_mode='border')
            concat = torch.cat([x_warp, x_cond], 1)
            flow = self.netRefine[i](concat)

            g_map = self.map_preds[i](concat)
            if i >= fea_num - self.use_num:
                upsample_tmp = 0.5**(fea_num - i)
                g_map = self.map_refine[i](
                    torch.cat([
                        g_map, last_gmap,
                        self.upsample(
                            localmap_list[i - (fea_num - self.use_num)],
                            upsample_tmp)
                    ], 1))
            elif last_gmap is not None:
                g_map = self.map_refine[i](torch.cat([g_map, last_gmap], 1))
            elif i == 0:
                g_map = self.map_refine[i](g_map)
            last_gmap = self.upsample(g_map, 2)
            gmap_all.append(last_gmap)

            delta_list.append(flow)
            flow = apply_offset(flow)
            flow = F.grid_sample(
                last_flow, flow, mode='bilinear', padding_mode='border')

            last_flow = F.interpolate(flow, scale_factor=2, mode='bilinear')
            last_flow_all.append(last_flow)
            cur_x = F.interpolate(
                x, scale_factor=0.5**(len(x_warps) - 1 - i), mode='bilinear')
            cur_x_warp = F.grid_sample(
                cur_x,
                last_flow.permute(0, 2, 3, 1),
                mode='bilinear',
                padding_mode='border')

            x_all.append(cur_x_warp)
            cur_x_edge = F.interpolate(
                x_edge,
                scale_factor=0.5**(len(x_warps) - 1 - i),
                mode='bilinear')
            cur_x_warp_edge = F.grid_sample(
                cur_x_edge,
                last_flow.permute(0, 2, 3, 1),
                mode='bilinear',
                padding_mode='zeros')
            x_edge_all.append(cur_x_warp_edge)
            flow_x, flow_y = torch.split(last_flow, 1, dim=1)
            delta_x = F.conv2d(flow_x, self.weight)
            delta_y = F.conv2d(flow_y, self.weight)
            delta_x_all.append(delta_x)
            delta_y_all.append(delta_y)

        x_warp = F.grid_sample(
            x,
            last_flow.permute(0, 2, 3, 1),
            mode='bilinear',
            padding_mode='border')
        return x_warp, last_flow, cond_fea_all, last_flow_all, delta_list, x_all, \
            x_edge_all, delta_x_all, delta_y_all, gmap_all

    def upsample(self, F, scale):
        """[2x nearest neighbor upsampling]
        Arguments:
            F {[torch tensor]} -- [tensor to be upsampled, (B, C, H, W)]
        """
        upsample = torch.nn.Upsample(scale_factor=scale, mode='nearest')
        return upsample(F)


class Warping(nn.Module):
    """initialize the try on warping model
    """

    def __init__(self, input_nc=32 + 3, use_num=3):
        super(Warping, self).__init__()
        num_filters = [64, 128, 256, 256, 256]
        self.num = len(num_filters) - 1
        self.cloth_input = nn.Conv2d(
            3, 32, kernel_size=3, padding=1, bias=False)
        self.person_input = nn.Conv2d(
            3, 32, kernel_size=3, padding=1, bias=False)
        self.c_input_refine = Conv2dBlock(
            64, 35, 3, 1, 1, activation='prelu', bias=False)
        self.p_input_refine = Conv2dBlock(
            64, 35, 3, 1, 1, activation='prelu', bias=False)

        self.src_features = FeatureEncoder(input_nc, num_filters)
        self.tar_features = FeatureEncoder(input_nc, num_filters)
        self.src_FPN = RefinePyramid(num_filters)
        self.tar_FPN = RefinePyramid(num_filters)
        self.global_flow = AFlowNet(len(num_filters), use_num=use_num)
        self.local_flow = LocalFlow(use_num=use_num, init_scale=5)

        self.softmax = nn.Softmax(dim=1)

        self.input_scale = 4

    def forward(self, warping_input):

        cloth = warping_input[0]
        person = warping_input[1]
        c_heatmap = warping_input[2]
        p_heatmap = warping_input[3]
        c_landmark = warping_input[4]
        p_landmark = warping_input[5]
        x_edge = warping_input[6]
        org_cloth = warping_input[7]
        device = warping_input[8]

        N, _, H, W = cloth.shape
        cf = self.cloth_input(cloth)
        pf = self.person_input(person)
        src_input = self.c_input_refine(torch.cat((cf, c_heatmap), dim=1))
        tar_input = self.p_input_refine(torch.cat((pf, p_heatmap), dim=1))

        src_en_fea = self.src_features(src_input)
        tar_en_fea = self.tar_features(tar_input)
        src_c_fea = self.src_FPN(src_en_fea)
        tar_c_fea = self.tar_FPN(tar_en_fea)
        src_fuse_fea = src_c_fea
        tar_fuse_fea = tar_c_fea

        localflow_list, localmap_list = self.local_flow(
            cloth, src_fuse_fea, tar_fuse_fea, c_landmark, p_landmark)
        global_flow_input = [
            cloth, x_edge, src_fuse_fea, tar_fuse_fea, localflow_list,
            localmap_list
        ]
        x_warp, last_flow, last_flow_all, flow_all, delta_list, x_all, x_edge_all, \
            delta_x_all, delta_y_all, gmap_all = self.global_flow(global_flow_input, device)

        local_warped_cloth_list = []
        for i in range(len(localflow_list)):
            localmap = self.softmax(localmap_list[i])
            warped_cloth = torch.zeros_like(cloth)
            for j in range(32):
                once_warped_cloth = F.grid_sample(
                    cloth,
                    localflow_list[i][:, j * 2:(j + 1) * 2, :, :].permute(
                        0, 2, 3, 1),
                    mode='bilinear',
                    padding_mode='border')
                warped_cloth += once_warped_cloth * localmap[:,
                                                             j, :, :].unsqueeze(
                                                                 dim=1)
            local_warped_cloth_list.append(warped_cloth)
        globalmap = self.softmax(gmap_all[-1])
        fuse_cloth = globalmap[:, 0, :, :].unsqueeze(dim=1) * x_warp + \
            globalmap[:, 1, :, :].unsqueeze(dim=1) * local_warped_cloth_list[-1]

        # global local fused
        # global
        _, _, h, w = org_cloth.shape
        up_last_flow = F.interpolate(last_flow, size=(h, w), mode='bicubic')
        up_warped_gcloth = F.grid_sample(
            org_cloth,
            up_last_flow.permute(0, 2, 3, 1),
            mode='bicubic',
            padding_mode='border')
        # local
        localmap = self.softmax(localmap_list[-1])
        up_map = F.interpolate(localmap, size=(h, w), mode='bicubic')
        up_flow = F.interpolate(
            localflow_list[-1], size=(h, w), mode='bicubic')
        up_warped_lcloth = torch.zeros_like(org_cloth)
        for j in range(32):
            once_up_flow = F.interpolate(
                up_flow[:, j * 2:(j + 1) * 2, :, :],
                size=(h, w),
                mode='bicubic')
            once_warped_cloth = F.grid_sample(
                org_cloth,
                once_up_flow.permute(0, 2, 3, 1),
                mode='bilinear',
                padding_mode='border')
            up_warped_lcloth += once_warped_cloth * up_map[:,
                                                           j, :, :].unsqueeze(
                                                               dim=1)
        up_globalmap = F.interpolate(globalmap, size=(h, w), mode='bicubic')
        up_fuse_cloth = up_globalmap[:, 0, :, :].unsqueeze(
            dim=1) * up_warped_gcloth + up_globalmap[:, 1, :, :].unsqueeze(
                dim=1) * up_warped_lcloth
        up_cloth = torch.cat(
            [up_warped_gcloth, up_warped_lcloth, up_fuse_cloth], 1)

        return x_warp, last_flow, last_flow_all, flow_all, delta_list, x_all, x_edge_all, \
            delta_x_all, delta_y_all, local_warped_cloth_list, fuse_cloth, gmap_all, up_cloth

    def warp(self, x, flo, device):
        """
        warp an image/tensor (im2) back to im1, according to the optical flow
        x: [B, C, H, W] (im2)
        flo: [B, 2, H, W] flow
        """
        import torch.autograd as autograd
        from torch.autograd import Variable
        B, C, H, W = x.size()
        # mesh grid
        xx = torch.arange(0, W).view(1, -1).repeat(H, 1)
        yy = torch.arange(0, H).view(-1, 1).repeat(1, W)
        xx = xx.view(1, 1, H, W).repeat(B, 1, 1, 1)
        yy = yy.view(1, 1, H, W).repeat(B, 1, 1, 1)
        grid = torch.cat((xx, yy), 1).float()

        if x.is_cuda:
            grid = grid.cuda()
        vgrid = Variable(grid) + flo

        # scale grid to [-1,1]
        vgrid[:,
              0, :, :] = 2.0 * vgrid[:, 0, :, :].clone() / max(W - 1, 1) - 1.0
        vgrid[:,
              1, :, :] = 2.0 * vgrid[:, 1, :, :].clone() / max(H - 1, 1) - 1.0

        vgrid = vgrid.permute(0, 2, 3, 1)
        output = nn.functional.grid_sample(x, vgrid)
        mask = torch.autograd.Variable(torch.ones(x.size())).to(device)
        mask = nn.functional.grid_sample(mask, vgrid)

        mask[mask < 0.9999] = 0
        mask[mask > 0] = 1

        return output * mask