Upload 5 files

networks.py

@@ -0,0 +1,541 @@
+# coding=utf-8
+import torch
+import torch.nn as nn
+from torch.nn import init
+from torchvision import models
+import os
+
+import numpy as np
+
+
+def weights_init_normal(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find('Linear') != -1:
+        init.normal(m.weight.data, 0.0, 0.02)
+    elif classname.find('BatchNorm2d') != -1:
+        init.normal_(m.weight.data, 1.0, 0.02)
+        init.constant_(m.bias.data, 0.0)
+
+
+def weights_init_xavier(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        init.xavier_normal_(m.weight.data, gain=0.02)
+    elif classname.find('Linear') != -1:
+        init.xavier_normal_(m.weight.data, gain=0.02)
+    elif classname.find('BatchNorm2d') != -1:
+        init.normal_(m.weight.data, 1.0, 0.02)
+        init.constant_(m.bias.data, 0.0)
+
+
+def weights_init_kaiming(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+    elif classname.find('Linear') != -1:
+        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+    elif classname.find('BatchNorm2d') != -1:
+        init.normal_(m.weight.data, 1.0, 0.02)
+        init.constant_(m.bias.data, 0.0)
+
+
+def init_weights(net, init_type='normal'):
+    print('initialization method [%s]' % init_type)
+    if init_type == 'normal':
+        net.apply(weights_init_normal)
+    elif init_type == 'xavier':
+        net.apply(weights_init_xavier)
+    elif init_type == 'kaiming':
+        net.apply(weights_init_kaiming)
+    else:
+        raise NotImplementedError(
+            'initialization method [%s] is not implemented' % init_type)
+
+
+class FeatureExtraction(nn.Module):
+    def __init__(self, input_nc, ngf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(FeatureExtraction, self).__init__()
+        downconv = nn.Conv2d(input_nc, ngf, kernel_size=4, stride=2, padding=1)
+        model = [downconv, nn.ReLU(True), norm_layer(ngf)]
+        for i in range(n_layers):
+            in_ngf = 2**i * ngf if 2**i * ngf < 512 else 512
+            out_ngf = 2**(i+1) * ngf if 2**i * ngf < 512 else 512
+            downconv = nn.Conv2d(
+                in_ngf, out_ngf, kernel_size=4, stride=2, padding=1)
+            model += [downconv, nn.ReLU(True)]
+            model += [norm_layer(out_ngf)]
+        model += [nn.Conv2d(512, 512, kernel_size=3,
+                            stride=1, padding=1), nn.ReLU(True)]
+        model += [norm_layer(512)]
+        model += [nn.Conv2d(512, 512, kernel_size=3,
+                            stride=1, padding=1), nn.ReLU(True)]
+
+        self.model = nn.Sequential(*model)
+        init_weights(self.model, init_type='normal')
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class FeatureL2Norm(torch.nn.Module):
+    def __init__(self):
+        super(FeatureL2Norm, self).__init__()
+
+    def forward(self, feature):
+        epsilon = 1e-6
+        norm = torch.pow(torch.sum(torch.pow(feature, 2), 1) +
+                         epsilon, 0.5).unsqueeze(1).expand_as(feature)
+        return torch.div(feature, norm)
+
+
+class FeatureCorrelation(nn.Module):
+    def __init__(self):
+        super(FeatureCorrelation, self).__init__()
+
+    def forward(self, feature_A, feature_B):
+        b, c, h, w = feature_A.size()
+        # reshape features for matrix multiplication
+        feature_A = feature_A.transpose(2, 3).contiguous().view(b, c, h*w)
+        feature_B = feature_B.view(b, c, h*w).transpose(1, 2)
+        # perform matrix mult.
+        feature_mul = torch.bmm(feature_B, feature_A)
+        correlation_tensor = feature_mul.view(
+            b, h, w, h*w).transpose(2, 3).transpose(1, 2)
+        return correlation_tensor
+
+
+class FeatureRegression(nn.Module):
+    def __init__(self, input_nc=512, output_dim=6, use_cuda=True):
+        super(FeatureRegression, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(input_nc, 512, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(512),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(512, 256, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+        )
+        self.linear = nn.Linear(64 * 4 * 3, output_dim)
+        self.tanh = nn.Tanh()
+        if use_cuda:
+            self.conv.cuda()
+            self.linear.cuda()
+            self.tanh.cuda()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = x.view(x.size(0), -1)
+        x = self.linear(x)
+        x = self.tanh(x)
+        return x
+
+
+class AffineGridGen(nn.Module):
+    def __init__(self, out_h=256, out_w=192, out_ch=3):
+        super(AffineGridGen, self).__init__()
+        self.out_h = out_h
+        self.out_w = out_w
+        self.out_ch = out_ch
+
+    def forward(self, theta):
+        theta = theta.contiguous()
+        batch_size = theta.size()[0]
+        out_size = torch.Size(
+            (batch_size, self.out_ch, self.out_h, self.out_w))
+        return F.affine_grid(theta, out_size)
+
+
+class TpsGridGen(nn.Module):
+    def __init__(self, out_h=256, out_w=192, use_regular_grid=True, grid_size=3, reg_factor=0, use_cuda=True):
+        super(TpsGridGen, self).__init__()
+        self.out_h, self.out_w = out_h, out_w
+        self.reg_factor = reg_factor
+        self.use_cuda = use_cuda
+
+        # create grid in numpy
+        self.grid = np.zeros([self.out_h, self.out_w, 3], dtype=np.float32)
+        # sampling grid with dim-0 coords (Y)
+        self.grid_X, self.grid_Y = np.meshgrid(
+            np.linspace(-1, 1, out_w), np.linspace(-1, 1, out_h))
+        # grid_X,grid_Y: size [1,H,W,1,1]
+        self.grid_X = torch.FloatTensor(self.grid_X).unsqueeze(0).unsqueeze(3)
+        self.grid_Y = torch.FloatTensor(self.grid_Y).unsqueeze(0).unsqueeze(3)
+        if use_cuda:
+            self.grid_X = self.grid_X.cuda()
+            self.grid_Y = self.grid_Y.cuda()
+
+        # initialize regular grid for control points P_i
+        if use_regular_grid:
+            axis_coords = np.linspace(-1, 1, grid_size)
+            self.N = grid_size*grid_size
+            P_Y, P_X = np.meshgrid(axis_coords, axis_coords)
+            P_X = np.reshape(P_X, (-1, 1))  # size (N,1)
+            P_Y = np.reshape(P_Y, (-1, 1))  # size (N,1)
+            P_X = torch.FloatTensor(P_X)
+            P_Y = torch.FloatTensor(P_Y)
+            self.P_X_base = P_X.clone()
+            self.P_Y_base = P_Y.clone()
+            self.Li = self.compute_L_inverse(P_X, P_Y).unsqueeze(0)
+            self.P_X = P_X.unsqueeze(2).unsqueeze(
+                3).unsqueeze(4).transpose(0, 4)
+            self.P_Y = P_Y.unsqueeze(2).unsqueeze(
+                3).unsqueeze(4).transpose(0, 4)
+            if use_cuda:
+                self.P_X = self.P_X.cuda()
+                self.P_Y = self.P_Y.cuda()
+                self.P_X_base = self.P_X_base.cuda()
+                self.P_Y_base = self.P_Y_base.cuda()
+
+    def forward(self, theta):
+        warped_grid = self.apply_transformation(
+            theta, torch.cat((self.grid_X, self.grid_Y), 3))
+
+        return warped_grid
+
+    def compute_L_inverse(self, X, Y):
+        N = X.size()[0]  # num of points (along dim 0)
+        # construct matrix K
+        Xmat = X.expand(N, N)
+        Ymat = Y.expand(N, N)
+        P_dist_squared = torch.pow(
+            Xmat-Xmat.transpose(0, 1), 2)+torch.pow(Ymat-Ymat.transpose(0, 1), 2)
+        # make diagonal 1 to avoid NaN in log computation
+        P_dist_squared[P_dist_squared == 0] = 1
+        K = torch.mul(P_dist_squared, torch.log(P_dist_squared))
+        # construct matrix L
+        O = torch.FloatTensor(N, 1).fill_(1)
+        Z = torch.FloatTensor(3, 3).fill_(0)
+        P = torch.cat((O, X, Y), 1)
+        L = torch.cat((torch.cat((K, P), 1), torch.cat(
+            (P.transpose(0, 1), Z), 1)), 0)
+        Li = torch.inverse(L)
+        if self.use_cuda:
+            Li = Li.cuda()
+        return Li
+
+    def apply_transformation(self, theta, points):
+        if theta.dim() == 2:
+            theta = theta.unsqueeze(2).unsqueeze(3)
+        # points should be in the [B,H,W,2] format,
+        # where points[:,:,:,0] are the X coords
+        # and points[:,:,:,1] are the Y coords
+
+        # input are the corresponding control points P_i
+        batch_size = theta.size()[0]
+        # split theta into point coordinates
+        Q_X = theta[:, :self.N, :, :].squeeze(3)
+        Q_Y = theta[:, self.N:, :, :].squeeze(3)
+        Q_X = Q_X + self.P_X_base.expand_as(Q_X)
+        Q_Y = Q_Y + self.P_Y_base.expand_as(Q_Y)
+
+        # get spatial dimensions of points
+        points_b = points.size()[0]
+        points_h = points.size()[1]
+        points_w = points.size()[2]
+
+        # repeat pre-defined control points along spatial dimensions of points to be transformed
+        P_X = self.P_X.expand((1, points_h, points_w, 1, self.N))
+        P_Y = self.P_Y.expand((1, points_h, points_w, 1, self.N))
+
+        # compute weigths for non-linear part
+        W_X = torch.bmm(self.Li[:, :self.N, :self.N].expand(
+            (batch_size, self.N, self.N)), Q_X)
+        W_Y = torch.bmm(self.Li[:, :self.N, :self.N].expand(
+            (batch_size, self.N, self.N)), Q_Y)
+        # reshape
+        # W_X,W,Y: size [B,H,W,1,N]
+        W_X = W_X.unsqueeze(3).unsqueeze(4).transpose(
+            1, 4).repeat(1, points_h, points_w, 1, 1)
+        W_Y = W_Y.unsqueeze(3).unsqueeze(4).transpose(
+            1, 4).repeat(1, points_h, points_w, 1, 1)
+        # compute weights for affine part
+        A_X = torch.bmm(self.Li[:, self.N:, :self.N].expand(
+            (batch_size, 3, self.N)), Q_X)
+        A_Y = torch.bmm(self.Li[:, self.N:, :self.N].expand(
+            (batch_size, 3, self.N)), Q_Y)
+        # reshape
+        # A_X,A,Y: size [B,H,W,1,3]
+        A_X = A_X.unsqueeze(3).unsqueeze(4).transpose(
+            1, 4).repeat(1, points_h, points_w, 1, 1)
+        A_Y = A_Y.unsqueeze(3).unsqueeze(4).transpose(
+            1, 4).repeat(1, points_h, points_w, 1, 1)
+
+        # compute distance P_i - (grid_X,grid_Y)
+        # grid is expanded in point dim 4, but not in batch dim 0, as points P_X,P_Y are fixed for all batch
+        points_X_for_summation = points[:, :, :, 0].unsqueeze(
+            3).unsqueeze(4).expand(points[:, :, :, 0].size()+(1, self.N))
+        points_Y_for_summation = points[:, :, :, 1].unsqueeze(
+            3).unsqueeze(4).expand(points[:, :, :, 1].size()+(1, self.N))
+
+        if points_b == 1:
+            delta_X = points_X_for_summation-P_X
+            delta_Y = points_Y_for_summation-P_Y
+        else:
+            # use expanded P_X,P_Y in batch dimension
+            delta_X = points_X_for_summation - \
+                P_X.expand_as(points_X_for_summation)
+            delta_Y = points_Y_for_summation - \
+                P_Y.expand_as(points_Y_for_summation)
+
+        dist_squared = torch.pow(delta_X, 2)+torch.pow(delta_Y, 2)
+        # U: size [1,H,W,1,N]
+        dist_squared[dist_squared == 0] = 1  # avoid NaN in log computation
+        U = torch.mul(dist_squared, torch.log(dist_squared))
+
+        # expand grid in batch dimension if necessary
+        points_X_batch = points[:, :, :, 0].unsqueeze(3)
+        points_Y_batch = points[:, :, :, 1].unsqueeze(3)
+        if points_b == 1:
+            points_X_batch = points_X_batch.expand(
+                (batch_size,)+points_X_batch.size()[1:])
+            points_Y_batch = points_Y_batch.expand(
+                (batch_size,)+points_Y_batch.size()[1:])
+
+        points_X_prime = A_X[:, :, :, :, 0] + \
+            torch.mul(A_X[:, :, :, :, 1], points_X_batch) + \
+            torch.mul(A_X[:, :, :, :, 2], points_Y_batch) + \
+            torch.sum(torch.mul(W_X, U.expand_as(W_X)), 4)
+
+        points_Y_prime = A_Y[:, :, :, :, 0] + \
+            torch.mul(A_Y[:, :, :, :, 1], points_X_batch) + \
+            torch.mul(A_Y[:, :, :, :, 2], points_Y_batch) + \
+            torch.sum(torch.mul(W_Y, U.expand_as(W_Y)), 4)
+
+        return torch.cat((points_X_prime, points_Y_prime), 3)
+
+# Defines the Unet generator.
+# |num_downs|: number of downsamplings in UNet. For example,
+# if |num_downs| == 7, image of size 128x128 will become of size 1x1
+# at the bottleneck
+
+
+class UnetGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetGenerator, self).__init__()
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+        for i in range(num_downs - 5):
+            unet_block = UnetSkipConnectionBlock(
+                ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(
+            ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(
+            output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+
+# Defines the submodule with skip connection.
+# X -------------------identity---------------------- X
+#   |-- downsampling -- |submodule| -- upsampling --|
+class UnetSkipConnectionBlock(nn.Module):
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        use_bias = norm_layer == nn.InstanceNorm2d
+
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                             stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upsample = nn.Upsample(scale_factor=2, mode='bilinear')
+            upconv = nn.Conv2d(inner_nc * 2, outer_nc,
+                               kernel_size=3, stride=1, padding=1, bias=use_bias)
+            down = [downconv]
+            up = [uprelu, upsample, upconv, upnorm]
+            model = down + [submodule] + up
+        elif innermost:
+            upsample = nn.Upsample(scale_factor=2, mode='bilinear')
+            upconv = nn.Conv2d(inner_nc, outer_nc, kernel_size=3,
+                               stride=1, padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upsample, upconv, upnorm]
+            model = down + up
+        else:
+            upsample = nn.Upsample(scale_factor=2, mode='bilinear')
+            upconv = nn.Conv2d(inner_nc*2, outer_nc, kernel_size=3,
+                               stride=1, padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upsample, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:
+            return torch.cat([x, self.model(x)], 1)
+
+
+class Vgg19(nn.Module):
+    def __init__(self, requires_grad=False):
+        super(Vgg19, self).__init__()
+        vgg_pretrained_features = models.vgg19(pretrained=True).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        for x in range(2):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(2, 7):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(7, 12):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(12, 21):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(21, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h_relu1 = self.slice1(X)
+        h_relu2 = self.slice2(h_relu1)
+        h_relu3 = self.slice3(h_relu2)
+        h_relu4 = self.slice4(h_relu3)
+        h_relu5 = self.slice5(h_relu4)
+        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
+        return out
+
+
+class VGGLoss(nn.Module):
+    def __init__(self, layids=None):
+        super(VGGLoss, self).__init__()
+        self.vgg = Vgg19()
+        self.vgg.cuda()
+        self.criterion = nn.L1Loss()
+        self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]
+        self.layids = layids
+
+    def forward(self, x, y):
+        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
+        loss = 0
+        if self.layids is None:
+            self.layids = list(range(len(x_vgg)))
+        for i in self.layids:
+            loss += self.weights[i] * \
+                self.criterion(x_vgg[i], y_vgg[i].detach())
+        return loss
+
+
+class DT(nn.Module):
+    def __init__(self):
+        super(DT, self).__init__()
+
+    def forward(self, x1, x2):
+        dt = torch.abs(x1 - x2)
+        return dt
+
+
+class DT2(nn.Module):
+    def __init__(self):
+        super(DT, self).__init__()
+
+    def forward(self, x1, y1, x2, y2):
+        dt = torch.sqrt(torch.mul(x1 - x2, x1 - x2) +
+                        torch.mul(y1 - y2, y1 - y2))
+        return dt
+
+
+class GicLoss(nn.Module):
+    def __init__(self, opt):
+        super(GicLoss, self).__init__()
+        self.dT = DT()
+        self.opt = opt
+
+    def forward(self, grid):
+        Gx = grid[:, :, :, 0]
+        Gy = grid[:, :, :, 1]
+        Gxcenter = Gx[:, 1:self.opt.fine_height - 1, 1:self.opt.fine_width - 1]
+        Gxup = Gx[:, 0:self.opt.fine_height - 2, 1:self.opt.fine_width - 1]
+        Gxdown = Gx[:, 2:self.opt.fine_height, 1:self.opt.fine_width - 1]
+        Gxleft = Gx[:, 1:self.opt.fine_height - 1, 0:self.opt.fine_width - 2]
+        Gxright = Gx[:, 1:self.opt.fine_height - 1, 2:self.opt.fine_width]
+
+        Gycenter = Gy[:, 1:self.opt.fine_height - 1, 1:self.opt.fine_width - 1]
+        Gyup = Gy[:, 0:self.opt.fine_height - 2, 1:self.opt.fine_width - 1]
+        Gydown = Gy[:, 2:self.opt.fine_height, 1:self.opt.fine_width - 1]
+        Gyleft = Gy[:, 1:self.opt.fine_height - 1, 0:self.opt.fine_width - 2]
+        Gyright = Gy[:, 1:self.opt.fine_height - 1, 2:self.opt.fine_width]
+
+        dtleft = self.dT(Gxleft, Gxcenter)
+        dtright = self.dT(Gxright, Gxcenter)
+        dtup = self.dT(Gyup, Gycenter)
+        dtdown = self.dT(Gydown, Gycenter)
+
+        return torch.sum(torch.abs(dtleft - dtright) + torch.abs(dtup - dtdown))
+
+
+class GMM(nn.Module):
+    """ Geometric Matching Module
+    """
+
+    def __init__(self, opt):
+        super(GMM, self).__init__()
+        self.extractionA = FeatureExtraction(
+            22, ngf=64, n_layers=3, norm_layer=nn.BatchNorm2d)
+        self.extractionB = FeatureExtraction(
+            1, ngf=64, n_layers=3, norm_layer=nn.BatchNorm2d)
+        self.l2norm = FeatureL2Norm()
+        self.correlation = FeatureCorrelation()
+        self.regression = FeatureRegression(
+            input_nc=192, output_dim=2*opt.grid_size**2, use_cuda=True)
+        self.gridGen = TpsGridGen(
+            opt.fine_height, opt.fine_width, use_cuda=True, grid_size=opt.grid_size)
+
+    def forward(self, inputA, inputB):
+        featureA = self.extractionA(inputA)
+        featureB = self.extractionB(inputB)
+        featureA = self.l2norm(featureA)
+        featureB = self.l2norm(featureB)
+        correlation = self.correlation(featureA, featureB)
+
+        theta = self.regression(correlation)
+        grid = self.gridGen(theta)
+        return grid, theta
+
+
+def save_checkpoint(model, save_path):
+    if not os.path.exists(os.path.dirname(save_path)):
+        os.makedirs(os.path.dirname(save_path))
+
+    torch.save(model.cpu().state_dict(), save_path)
+    model.cuda()
+
+
+def load_checkpoint(model, checkpoint_path):
+    if not os.path.exists(checkpoint_path):
+        return
+    model.load_state_dict(torch.load(checkpoint_path))
+    model.cuda()

requirements (2).txt

@@ -0,0 +1,6 @@
+torch>=1.10
+torchvision>=0.11
+tensorboardX
+pillow
+numpy
+opencv-contrib-python

test.py

@@ -0,0 +1,226 @@
+# coding=utf-8
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import argparse
+import os
+import time
+from cp_dataset import CPDataset, CPDataLoader
+from networks import GMM, UnetGenerator, load_checkpoint
+
+from tensorboardX import SummaryWriter
+from visualization import board_add_image, board_add_images, save_images
+
+
+def get_opt():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--name", default="GMM")
+    # parser.add_argument("--name", default="TOM")
+
+    parser.add_argument("--gpu_ids", default="")
+    parser.add_argument('-j', '--workers', type=int, default=1)
+    parser.add_argument('-b', '--batch-size', type=int, default=4)
+
+    parser.add_argument("--dataroot", default="data")
+
+    # parser.add_argument("--datamode", default="train")
+    parser.add_argument("--datamode", default="test")
+
+    parser.add_argument("--stage", default="GMM")
+    # parser.add_argument("--stage", default="TOM")
+
+    # parser.add_argument("--data_list", default="train_pairs.txt")
+    parser.add_argument("--data_list", default="test_pairs.txt")
+    # parser.add_argument("--data_list", default="test_pairs_same.txt")
+
+    parser.add_argument("--fine_width", type=int, default=192)
+    parser.add_argument("--fine_height", type=int, default=256)
+    parser.add_argument("--radius", type=int, default=5)
+    parser.add_argument("--grid_size", type=int, default=5)
+
+    parser.add_argument('--tensorboard_dir', type=str,
+                        default='tensorboard', help='save tensorboard infos')
+
+    parser.add_argument('--result_dir', type=str,
+                        default='result', help='save result infos')
+
+    parser.add_argument('--checkpoint', type=str, default='checkpoints/GMM/gmm_final.pth', help='model checkpoint for test')
+    # parser.add_argument('--checkpoint', type=str, default='checkpoints/TOM/tom_final.pth', help='model checkpoint for test')
+
+    parser.add_argument("--display_count", type=int, default=1)
+    parser.add_argument("--shuffle", action='store_true',
+                        help='shuffle input data')
+
+    opt = parser.parse_args()
+    return opt
+
+
+def test_gmm(opt, test_loader, model, board):
+    model.cuda()
+    model.eval()
+
+    base_name = os.path.basename(opt.checkpoint)
+    name = opt.name
+    save_dir = os.path.join(opt.result_dir, name, opt.datamode)
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    warp_cloth_dir = os.path.join(save_dir, 'warp-cloth')
+    if not os.path.exists(warp_cloth_dir):
+        os.makedirs(warp_cloth_dir)
+    warp_mask_dir = os.path.join(save_dir, 'warp-mask')
+    if not os.path.exists(warp_mask_dir):
+        os.makedirs(warp_mask_dir)
+    result_dir1 = os.path.join(save_dir, 'result_dir')
+    if not os.path.exists(result_dir1):
+        os.makedirs(result_dir1)
+    overlayed_TPS_dir = os.path.join(save_dir, 'overlayed_TPS')
+    if not os.path.exists(overlayed_TPS_dir):
+        os.makedirs(overlayed_TPS_dir)
+    warped_grid_dir = os.path.join(save_dir, 'warped_grid')
+    if not os.path.exists(warped_grid_dir):
+        os.makedirs(warped_grid_dir)
+    for step, inputs in enumerate(test_loader.data_loader):
+        iter_start_time = time.time()
+
+        c_names = inputs['c_name']
+        im_names = inputs['im_name']
+        im = inputs['image'].cuda()
+        im_pose = inputs['pose_image'].cuda()
+        im_h = inputs['head'].cuda()
+        shape = inputs['shape'].cuda()
+        agnostic = inputs['agnostic'].cuda()
+        c = inputs['cloth'].cuda()
+        cm = inputs['cloth_mask'].cuda()
+        im_c = inputs['parse_cloth'].cuda()
+        im_g = inputs['grid_image'].cuda()
+        shape_ori = inputs['shape_ori']  # original body shape without blurring
+
+        grid, theta = model(agnostic, cm)
+        warped_cloth = F.grid_sample(c, grid, padding_mode='border')
+        warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
+        warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
+        overlay = 0.7 * warped_cloth + 0.3 * im
+
+        visuals = [[im_h, shape, im_pose],
+                   [c, warped_cloth, im_c],
+                   [warped_grid, (warped_cloth+im)*0.5, im]]
+
+        # save_images(warped_cloth, c_names, warp_cloth_dir)
+        # save_images(warped_mask*2-1, c_names, warp_mask_dir)
+        save_images(warped_cloth, im_names, warp_cloth_dir)
+        save_images(warped_mask * 2 - 1, im_names, warp_mask_dir)
+        save_images(shape_ori.cuda() * 0.2 + warped_cloth *
+                    0.8, im_names, result_dir1)
+        save_images(warped_grid, im_names, warped_grid_dir)
+        save_images(overlay, im_names, overlayed_TPS_dir)
+
+        if (step+1) % opt.display_count == 0:
+            board_add_images(board, 'combine', visuals, step+1)
+            t = time.time() - iter_start_time
+            print('step: %8d, time: %.3f' % (step+1, t), flush=True)
+
+
+def test_tom(opt, test_loader, model, board):
+    model.cuda()
+    model.eval()
+
+    base_name = os.path.basename(opt.checkpoint)
+    # save_dir = os.path.join(opt.result_dir, base_name, opt.datamode)
+    save_dir = os.path.join(opt.result_dir, opt.name, opt.datamode)
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    try_on_dir = os.path.join(save_dir, 'try-on')
+    if not os.path.exists(try_on_dir):
+        os.makedirs(try_on_dir)
+    p_rendered_dir = os.path.join(save_dir, 'p_rendered')
+    if not os.path.exists(p_rendered_dir):
+        os.makedirs(p_rendered_dir)
+    m_composite_dir = os.path.join(save_dir, 'm_composite')
+    if not os.path.exists(m_composite_dir):
+        os.makedirs(m_composite_dir)
+    im_pose_dir = os.path.join(save_dir, 'im_pose')
+    if not os.path.exists(im_pose_dir):
+        os.makedirs(im_pose_dir)
+    shape_dir = os.path.join(save_dir, 'shape')
+    if not os.path.exists(shape_dir):
+        os.makedirs(shape_dir)
+    im_h_dir = os.path.join(save_dir, 'im_h')
+    if not os.path.exists(im_h_dir):
+        os.makedirs(im_h_dir)  # for test data
+
+    print('Dataset size: %05d!' % (len(test_loader.dataset)), flush=True)
+    for step, inputs in enumerate(test_loader.data_loader):
+        iter_start_time = time.time()
+
+        im_names = inputs['im_name']
+        im = inputs['image'].cuda()
+        im_pose = inputs['pose_image']
+        im_h = inputs['head']
+        shape = inputs['shape']
+
+        agnostic = inputs['agnostic'].cuda()
+        c = inputs['cloth'].cuda()
+        cm = inputs['cloth_mask'].cuda()
+
+        # outputs = model(torch.cat([agnostic, c], 1))  # CP-VTON
+        outputs = model(torch.cat([agnostic, c, cm], 1))  # CP-VTON+
+        p_rendered, m_composite = torch.split(outputs, 3, 1)
+        p_rendered = F.tanh(p_rendered)
+        m_composite = F.sigmoid(m_composite)
+        p_tryon = c * m_composite + p_rendered * (1 - m_composite)
+
+        visuals = [[im_h, shape, im_pose],
+                   [c, 2*cm-1, m_composite],
+                   [p_rendered, p_tryon, im]]
+
+        save_images(p_tryon, im_names, try_on_dir)
+        save_images(im_h, im_names, im_h_dir)
+        save_images(shape, im_names, shape_dir)
+        save_images(im_pose, im_names, im_pose_dir)
+        save_images(m_composite, im_names, m_composite_dir)
+        save_images(p_rendered, im_names, p_rendered_dir)  # For test data
+
+        if (step+1) % opt.display_count == 0:
+            board_add_images(board, 'combine', visuals, step+1)
+            t = time.time() - iter_start_time
+            print('step: %8d, time: %.3f' % (step+1, t), flush=True)
+
+
+def main():
+    opt = get_opt()
+    print(opt)
+    print("Start to test stage: %s, named: %s!" % (opt.stage, opt.name))
+
+    # create dataset
+    test_dataset = CPDataset(opt)
+
+    # create dataloader
+    test_loader = CPDataLoader(opt, test_dataset)
+
+    # visualization
+    if not os.path.exists(opt.tensorboard_dir):
+        os.makedirs(opt.tensorboard_dir)
+    board = SummaryWriter(logdir=os.path.join(opt.tensorboard_dir, opt.name))
+
+    # create model & test
+    if opt.stage == 'GMM':
+        model = GMM(opt)
+        load_checkpoint(model, opt.checkpoint)
+        with torch.no_grad():
+            test_gmm(opt, test_loader, model, board)
+    elif opt.stage == 'TOM':
+        # model = UnetGenerator(25, 4, 6, ngf=64, norm_layer=nn.InstanceNorm2d)  # CP-VTON
+        model = UnetGenerator(26, 4, 6, ngf=64, norm_layer=nn.InstanceNorm2d)  # CP-VTON+
+        load_checkpoint(model, opt.checkpoint)
+        with torch.no_grad():
+            test_tom(opt, test_loader, model, board)
+    else:
+        raise NotImplementedError('Model [%s] is not implemented' % opt.stage)
+
+    print('Finished test %s, named: %s!' % (opt.stage, opt.name))
+
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,232 @@
+# coding=utf-8
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import argparse
+import os
+import time
+from cp_dataset import CPDataset, CPDataLoader
+from networks import GicLoss, GMM, UnetGenerator, VGGLoss, load_checkpoint, save_checkpoint
+
+from tensorboardX import SummaryWriter
+from visualization import board_add_image, board_add_images
+
+
+def get_opt():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--name", default="GMM")
+    # parser.add_argument("--name", default="TOM")
+
+    parser.add_argument("--gpu_ids", default="")
+    parser.add_argument('-j', '--workers', type=int, default=1)
+    parser.add_argument('-b', '--batch-size', type=int, default=4)
+
+    parser.add_argument("--dataroot", default="data")
+
+    parser.add_argument("--datamode", default="train")
+
+    parser.add_argument("--stage", default="GMM")
+    # parser.add_argument("--stage", default="TOM")
+
+    parser.add_argument("--data_list", default="train_pairs.txt")
+
+    parser.add_argument("--fine_width", type=int, default=192)
+    parser.add_argument("--fine_height", type=int, default=256)
+    parser.add_argument("--radius", type=int, default=5)
+    parser.add_argument("--grid_size", type=int, default=5)
+    parser.add_argument('--lr', type=float, default=0.0001,
+                        help='initial learning rate for adam')
+    parser.add_argument('--tensorboard_dir', type=str,
+                        default='tensorboard', help='save tensorboard infos')
+    parser.add_argument('--checkpoint_dir', type=str,
+                        default='checkpoints', help='save checkpoint infos')
+    parser.add_argument('--checkpoint', type=str, default='',
+                        help='model checkpoint for initialization')
+    parser.add_argument("--display_count", type=int, default=20)
+    parser.add_argument("--save_count", type=int, default=5000)
+    parser.add_argument("--keep_step", type=int, default=100000)
+    parser.add_argument("--decay_step", type=int, default=100000)
+    parser.add_argument("--shuffle", action='store_true',
+                        help='shuffle input data')
+
+    opt = parser.parse_args()
+    return opt
+
+
+def train_gmm(opt, train_loader, model, board):
+    model.cuda()
+    model.train()
+
+    # criterion
+    criterionL1 = nn.L1Loss()
+    gicloss = GicLoss(opt)
+    # optimizer
+    optimizer = torch.optim.Adam(
+        model.parameters(), lr=opt.lr, betas=(0.5, 0.999))
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda step: 1.0 -
+                                                  max(0, step - opt.keep_step) / float(opt.decay_step + 1))
+
+    for step in range(opt.keep_step + opt.decay_step):
+        iter_start_time = time.time()
+        inputs = train_loader.next_batch()
+
+        im = inputs['image'].cuda()
+        im_pose = inputs['pose_image'].cuda()
+        im_h = inputs['head'].cuda()
+        shape = inputs['shape'].cuda()
+        agnostic = inputs['agnostic'].cuda()
+        c = inputs['cloth'].cuda()
+        cm = inputs['cloth_mask'].cuda()
+        im_c = inputs['parse_cloth'].cuda()
+        im_g = inputs['grid_image'].cuda()
+
+        grid, theta = model(agnostic, cm)    # can be added c too for new training
+        warped_cloth = F.grid_sample(c, grid, padding_mode='border')
+        warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
+        warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
+
+        visuals = [[im_h, shape, im_pose],
+                   [c, warped_cloth, im_c],
+                   [warped_grid, (warped_cloth+im)*0.5, im]]
+
+        # loss for warped cloth
+        Lwarp = criterionL1(warped_cloth, im_c)    # changing to previous code as it corresponds to the working code
+        # Actual loss function as in the paper given below (comment out previous line and uncomment below to train as per the paper)
+        # Lwarp = criterionL1(warped_mask, cm)    # loss for warped mask thanks @xuxiaochun025 for fixing the git code.
+        
+        # grid regularization loss
+        Lgic = gicloss(grid)
+        # 200x200 = 40.000 * 0.001
+        Lgic = Lgic / (grid.shape[0] * grid.shape[1] * grid.shape[2])
+
+        loss = Lwarp + 40 * Lgic    # total GMM loss
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        if (step+1) % opt.display_count == 0:
+            board_add_images(board, 'combine', visuals, step+1)
+            board.add_scalar('loss', loss.item(), step+1)
+            board.add_scalar('40*Lgic', (40*Lgic).item(), step+1)
+            board.add_scalar('Lwarp', Lwarp.item(), step+1)
+            t = time.time() - iter_start_time
+            print('step: %8d, time: %.3f, loss: %4f, (40*Lgic): %.8f, Lwarp: %.6f' %
+                  (step+1, t, loss.item(), (40*Lgic).item(), Lwarp.item()), flush=True)
+
+        if (step+1) % opt.save_count == 0:
+            save_checkpoint(model, os.path.join(
+                opt.checkpoint_dir, opt.name, 'step_%06d.pth' % (step+1)))
+
+
+def train_tom(opt, train_loader, model, board):
+    model.cuda()
+    model.train()
+
+    # criterion
+    criterionL1 = nn.L1Loss()
+    criterionVGG = VGGLoss()
+    criterionMask = nn.L1Loss()
+
+    # optimizer
+    optimizer = torch.optim.Adam(
+        model.parameters(), lr=opt.lr, betas=(0.5, 0.999))
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda step: 1.0 -
+                                                  max(0, step - opt.keep_step) / float(opt.decay_step + 1))
+
+    for step in range(opt.keep_step + opt.decay_step):
+        iter_start_time = time.time()
+        inputs = train_loader.next_batch()
+
+        im = inputs['image'].cuda()
+        im_pose = inputs['pose_image']
+        im_h = inputs['head']
+        shape = inputs['shape']
+
+        agnostic = inputs['agnostic'].cuda()
+        c = inputs['cloth'].cuda()
+        cm = inputs['cloth_mask'].cuda()
+        pcm = inputs['parse_cloth_mask'].cuda()
+
+        # outputs = model(torch.cat([agnostic, c], 1))  # CP-VTON
+        outputs = model(torch.cat([agnostic, c, cm], 1))  # CP-VTON+
+        p_rendered, m_composite = torch.split(outputs, 3, 1)
+        p_rendered = F.tanh(p_rendered)
+        m_composite = F.sigmoid(m_composite)
+        p_tryon = c * m_composite + p_rendered * (1 - m_composite)
+
+        """visuals = [[im_h, shape, im_pose],
+                   [c, cm*2-1, m_composite*2-1],
+                   [p_rendered, p_tryon, im]]"""  # CP-VTON
+
+        visuals = [[im_h, shape, im_pose],
+                   [c, pcm*2-1, m_composite*2-1],
+                   [p_rendered, p_tryon, im]]  # CP-VTON+
+
+        loss_l1 = criterionL1(p_tryon, im)
+        loss_vgg = criterionVGG(p_tryon, im)
+        # loss_mask = criterionMask(m_composite, cm)  # CP-VTON
+        loss_mask = criterionMask(m_composite, pcm)  # CP-VTON+
+        loss = loss_l1 + loss_vgg + loss_mask
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        if (step+1) % opt.display_count == 0:
+            board_add_images(board, 'combine', visuals, step+1)
+            board.add_scalar('metric', loss.item(), step+1)
+            board.add_scalar('L1', loss_l1.item(), step+1)
+            board.add_scalar('VGG', loss_vgg.item(), step+1)
+            board.add_scalar('MaskL1', loss_mask.item(), step+1)
+            t = time.time() - iter_start_time
+            print('step: %8d, time: %.3f, loss: %.4f, l1: %.4f, vgg: %.4f, mask: %.4f'
+                  % (step+1, t, loss.item(), loss_l1.item(),
+                     loss_vgg.item(), loss_mask.item()), flush=True)
+
+        if (step+1) % opt.save_count == 0:
+            save_checkpoint(model, os.path.join(
+                opt.checkpoint_dir, opt.name, 'step_%06d.pth' % (step+1)))
+
+
+def main():
+    opt = get_opt()
+    print(opt)
+    print("Start to train stage: %s, named: %s!" % (opt.stage, opt.name))
+
+    # create dataset
+    train_dataset = CPDataset(opt)
+
+    # create dataloader
+    train_loader = CPDataLoader(opt, train_dataset)
+
+    # visualization
+    if not os.path.exists(opt.tensorboard_dir):
+        os.makedirs(opt.tensorboard_dir)
+    board = SummaryWriter(logdir=os.path.join(opt.tensorboard_dir, opt.name))
+
+    # create model & train & save the final checkpoint
+    if opt.stage == 'GMM':
+        model = GMM(opt)
+        if not opt.checkpoint == '' and os.path.exists(opt.checkpoint):
+            load_checkpoint(model, opt.checkpoint)
+        train_gmm(opt, train_loader, model, board)
+        save_checkpoint(model, os.path.join(
+            opt.checkpoint_dir, opt.name, 'gmm_final.pth'))
+    elif opt.stage == 'TOM':
+        # model = UnetGenerator(25, 4, 6, ngf=64, norm_layer=nn.InstanceNorm2d)  # CP-VTON
+        model = UnetGenerator(
+            26, 4, 6, ngf=64, norm_layer=nn.InstanceNorm2d)  # CP-VTON+
+        if not opt.checkpoint == '' and os.path.exists(opt.checkpoint):
+            load_checkpoint(model, opt.checkpoint)
+        train_tom(opt, train_loader, model, board)
+        save_checkpoint(model, os.path.join(
+            opt.checkpoint_dir, opt.name, 'tom_final.pth'))
+    else:
+        raise NotImplementedError('Model [%s] is not implemented' % opt.stage)
+
+    print('Finished training %s, named: %s!' % (opt.stage, opt.name))
+
+
+if __name__ == "__main__":
+    main()

visualization.py

@@ -0,0 +1,64 @@
+from tensorboardX import SummaryWriter
+import torch
+from PIL import Image
+import os
+
+
+def tensor_for_board(img_tensor):
+    # map into [0,1]
+    tensor = (img_tensor.clone()+1) * 0.5
+    tensor.cpu().clamp(0, 1)
+
+    if tensor.size(1) == 1:
+        tensor = tensor.repeat(1, 3, 1, 1)
+
+    return tensor
+
+
+def tensor_list_for_board(img_tensors_list):
+    grid_h = len(img_tensors_list)
+    grid_w = max(len(img_tensors) for img_tensors in img_tensors_list)
+
+    batch_size, channel, height, width = tensor_for_board(
+        img_tensors_list[0][0]).size()
+    canvas_h = grid_h * height
+    canvas_w = grid_w * width
+    canvas = torch.FloatTensor(
+        batch_size, channel, canvas_h, canvas_w).fill_(0.5)
+    for i, img_tensors in enumerate(img_tensors_list):
+        for j, img_tensor in enumerate(img_tensors):
+            offset_h = i * height
+            offset_w = j * width
+            tensor = tensor_for_board(img_tensor)
+            canvas[:, :, offset_h: offset_h + height,
+                   offset_w: offset_w + width].copy_(tensor)
+
+    return canvas
+
+
+def board_add_image(board, tag_name, img_tensor, step_count):
+    tensor = tensor_for_board(img_tensor)
+
+    for i, img in enumerate(tensor):
+        board.add_image('%s/%03d' % (tag_name, i), img, step_count)
+
+
+def board_add_images(board, tag_name, img_tensors_list, step_count):
+    tensor = tensor_list_for_board(img_tensors_list)
+
+    for i, img in enumerate(tensor):
+        board.add_image('%s/%03d' % (tag_name, i), img, step_count)
+
+
+def save_images(img_tensors, img_names, save_dir):
+    for img_tensor, img_name in zip(img_tensors, img_names):
+        tensor = (img_tensor.clone()+1)*0.5 * 255
+        tensor = tensor.cpu().clamp(0, 255)
+
+        array = tensor.numpy().astype('uint8')
+        if array.shape[0] == 1:
+            array = array.squeeze(0)
+        elif array.shape[0] == 3:
+            array = array.swapaxes(0, 1).swapaxes(1, 2)
+
+        Image.fromarray(array).save(os.path.join(save_dir, img_name))