<jupyter_start><jupyter_text><jupyter_code>from google.colab import drive
drive.mount('/content/drive')
!cp /content/drive/Shareddrives/Virtuon/Pytorch/cp-vton-plus.zip /content/
!unzip -qq cp-vton-plus.zip -d /content/
# !pip install tensorboardX
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw
import torch
import torch.nn as nn
import torch.utils.data as data
import torchvision.transforms as transforms
import os.path as osp
import os
import json
import torch.nn.functional as F
import time
from torch.utils import tensorboard

%matplotlib inline
class CPDataset(data.Dataset):
    def __init__(self, stage='GMM', all_root="cp-vton-plus", data_path = "data", mode="train", radius=5, img_height=256, img_width=192):
        super(CPDataset, self).__init__()

        self.root = all_root

        self.data_root = osp.join(all_root,data_path)

        self.datamode = mode

        self.stage = stage

        self.data_list = "".join([mode, "_pairs.txt"])

        self.fine_height = img_height

        self.fine_width = img_width

        self.radius = radius

        self.data_path = osp.join(all_root,data_path, mode)
        
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        
        self.transform_1 = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5), (0.5))
        ])

        self.transform_2 = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5), (0.5, 0.5))
        ])

        self.transform_3 = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        im_names = []
        c_names = []

        with open(osp.join(self.data_root, self.data_list), 'r') as f:
            for line in f.readlines():
                im_name, c_name = line.strip().split()
                im_names.append(im_name)
                c_names.append(c_name)

        self.im_names = im_names
        self.c_names = c_names

    def name(self):
        return "CPDataset"
    
    def __getitem__(self, index):
        c_name = self.c_names[index]
        im_name = self.im_names[index]
        if self.stage == "GMM":
            c = Image.open(osp.join(self.data_path, 'cloth', c_name))
            cm = Image.open(osp.join(self.data_path, 'cloth-mask', c_name)).convert('L')
        else:
            c = Image.open(osp.join(self.data_path, 'warp-cloth', im_name))
            cm = Image.open(osp.join(self.data_path, 'warp-mask', im_name)).convert('L')
        
        c = self.transform(c)
        cm_array = np.array(cm)
        cm_array = (cm_array >= 128).astype(np.float32)
        cm = torch.from_numpy(cm_array)
        cm.unsqueeze_(0)

        # person image
        im = Image.open(osp.join(self.data_path, 'image', im_name))
        im = self.transform(im)

        
        # LIP labels
        
        # [(0, 0, 0),    # 0=Background
        #  (128, 0, 0),  # 1=Hat
        #  (255, 0, 0),  # 2=Hair
        #  (0, 85, 0),   # 3=Glove
        #  (170, 0, 51),  # 4=SunGlasses
        #  (255, 85, 0),  # 5=UpperClothes
        #  (0, 0, 85),     # 6=Dress
        #  (0, 119, 221),  # 7=Coat
        #  (85, 85, 0),    # 8=Socks
        #  (0, 85, 85),    # 9=Pants
        #  (85, 51, 0),    # 10=Jumpsuits
        #  (52, 86, 128),  # 11=Scarf
        #  (0, 128, 0),    # 12=Skirt
        #  (0, 0, 255),    # 13=Face
        #  (51, 170, 221),  # 14=LeftArm
        #  (0, 255, 255),   # 15=RightArm
        #  (85, 255, 170),  # 16=LeftLeg
        #  (170, 255, 85),  # 17=RightLeg
        #  (255, 255, 0),   # 18=LeftShoe
        #  (255, 170, 0)    # 19=RightShoe
        #  (170, 170, 50)   # 20=Skin/Neck/Chest (Newly added after running dataset_neck_skin_correction.py)
        #  ]
         
        # load parsing image
        parse_name = im_name.replace('.jpg', '.png')
        im_parse = Image.open(osp.join(self.data_path, 'image-parse-new',parse_name)).convert('L')
        parse_array = np.array(im_parse)

        im_mask = Image.open(osp.join(self.data_path, 'image-mask', parse_name)).convert('L')
        mask_array = np.array(im_mask)

        parse_shape = (mask_array > 0).astype(np.float32)

        if self.stage == 'GMM':
            parse_head = (parse_array == 1).astype(np.float32) + (parse_array == 4).astype(np.float32) + (parse_array == 13).astype(np.float32)

        else:
            parse_head = (parse_array == 1).astype(np.float32) + (parse_array == 2).astype(np.float32) + (parse_array == 4).astype(np.float32) + (parse_array == 9).astype(np.float32) + (parse_array == 12).astype(np.float32) + (parse_array == 13).astype(np.float32) + (parse_array == 16).astype(np.float32) + (parse_array == 17).astype(np.float32)  
            
        parse_cloth = (parse_array == 5).astype(np.float32) + (parse_array == 6).astype(np.float32) + (parse_array == 7).astype(np.float32)

        parse_shape_ori = Image.fromarray((parse_shape*255).astype(np.uint8))

        parse_shape = parse_shape_ori.resize((self.fine_width//16, self.fine_height//16), Image.BILINEAR)

        parse_shape = parse_shape.resize((self.fine_width, self.fine_height), Image.BILINEAR)
        
        parse_shape_ori = parse_shape_ori.resize((self.fine_width, self.fine_height), Image.BILINEAR)
        
        shape_ori = self.transform_1(parse_shape_ori)

        shape = self.transform_1(parse_shape)

        phead = torch.from_numpy(parse_head)

        pcm = torch.from_numpy(parse_cloth)

        # Upper Cloth
        im_c = im*pcm + (1 - pcm)
        im_h = im*phead + (1-phead)

        # load pose points
        pose_name = im_name.replace('.jpg', '_keypoints.json')
        with open(osp.join(self.data_path, 'pose', pose_name), 'r') as f:
            pose_label = json.load(f)
            pose_data = pose_label['people'][0]['pose_keypoints']
            pose_data = np.array(pose_data)
            pose_data = pose_data.reshape([-1,3])
        
        point_num = pose_data.shape[0]
        pose_map = torch.zeros(point_num, self.fine_height, self.fine_width)
        
        r = self.radius
        
        im_pose = Image.new('L', (self.fine_width, self.fine_height))
        pose_draw = ImageDraw.Draw(im_pose)

        for i in range(point_num):
            one_map = Image.new('L', (self.fine_width, self.fine_height))
            draw = ImageDraw.Draw(one_map)
            pointx = pose_data[i, 0]
            pointy = pose_data[i, 1]

            if pointx > 1 and pointy > 1:
                draw.rectangle((pointx - r, pointy - r, pointx + r, pointy + r), 'white', 'white')
                pose_draw.rectangle((pointx - r, pointy - r, pointx + r, pointy + r), 'white', 'white')

            one_map = self.transform_1(one_map)
            pose_map[i] = one_map[0]

        im_pose = self.transform_1(im_pose)

        agnostic = torch.cat([shape, im_h, pose_map], 0)

        if self.stage == 'GMM':
            im_g = Image.open(osp.join(self.root, 'grid.png'))
            im_g = self.transform(im_g)
        else:
            im_g = ''
        
        pcm.unsqueeze_(0)
        
        result = {
            'c_name': c_name,
            'im_name': im_name,
            'cloth': c,
            'cloth_mask': cm,
            'image': im,
            'agnostic': agnostic,
            'parse_cloth': im_c,
            'shape': shape,
            'head': im_h,
            'pose_image': im_pose,
            'grid_image': im_g,
            'parse_cloth_mask': pcm,
            'shape_ori': shape_ori,
        }

        return result
    def __len__(self):
        return len(self.im_names)


class CPDataLoader(object):
    def __init__(self, dataset, shuffle=True, batch=4, workers=4):
        super(CPDataLoader, self).__init__()

        if shuffle:
            train_sampler = torch.utils.data.sampler.RandomSampler(dataset)
        else:
            train_sampler = None
        
        self.data_loader = torch.utils.data.DataLoader(
            dataset, batch_size=batch, shuffle=(train_sampler is None),
            num_workers=workers, pin_memory=True, sampler=train_sampler
        )
        self.dataset = dataset
        self.data_iter = self.data_loader.__iter__()

    def next_batch(self):
        try:
            batch = self.data_iter.__next__()
        except StopIteration:
            self.data_iter = self.data_loader.__iter__()
            batch = self.data_iter.__next__()
        
        return batch

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 FeatureExtraction(nn.Module):
    def __init__(self, input_nc, ngf=64, n_layers=3, 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), nn.BatchNorm2d(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.append(downconv)
            model.append(nn.ReLU(True))
            model.append(nn.BatchNorm2d(out_ngf))
        
        model.append(nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1))
        model.append(nn.ReLU(True))
        model.append(nn.BatchNorm2d(512))
        model.append(nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1))
        model.append(nn.ReLU(True))

        self.model = nn.Sequential(*model)
        
    def forward(self, x):
        return self.model(x)



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()
        
        feature_A = feature_A.transpose(2,3).contiguous().view(b, c, h*w)
        feature_B = feature_B.contiguous().view(b, c, h*w).transpose(1,2)

        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
        # return feature_mul

class FeatureRegression(nn.Module):
    def __init__(self, input_nc=512, output_dim=50, 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(True),
            nn.Conv2d(512, 256, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(True)
        )
        self.linear = nn.Linear(64 * 4 * 3, output_dim)
        self.tanh = nn.Tanh()
        if use_cuda:
            self.conv.cuda()
            self.tanh.cuda()
            self.linear.cuda()
        
    def forward(self, x):
        x = self.conv(x)
        # x = x.view(x.size(0), -1)
        x = x.reshape(x.shape[0], -1)
        x = self.linear(x)
        x = self.tanh(x)
        return x
class TpsGridGen(nn.Module):
    def __init__(self, out_h=256, out_w=192, use_regular_grid=True, grid_size=5, 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)
class GMM(nn.Module):
    def __init__(self, grid_size = 5, fine_height=256, fine_width=192):
        super(GMM, self).__init__()
        self.extractionA = FeatureExtraction(22, ngf=64, n_layers=3)
        self.extractionB = FeatureExtraction(1, ngf=64, n_layers=3)
        self.l2norm = FeatureL2Norm()
        self.correlation = FeatureCorrelation()
        self.regression = FeatureRegression(input_nc=192, output_dim=2*grid_size**2, use_cuda=True)
        self.gridGen = TpsGridGen(fine_height, fine_width, use_cuda=True, grid_size=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.cuda(), featureB.cuda())

        theta = self.regression(correlation)
        grid = self.gridGen(theta)
        return grid, theta

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(DT2, 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, fine_height=256, fine_width=192):
        super(GicLoss, self).__init__()

        self.dT = DT()

        self.fine_height = fine_height
        self.fine_width = fine_width

    def forward(self, grid):
        Gx = grid[:,:,:,0]
        Gy = grid[:,:,:,1]
        
        Gxcenter = Gx[:, 1:self.fine_height - 1, 1:self.fine_width - 1]
        # Gxup = Gx[:, 0:self.fine_height - 2, 1:self.fine_width - 1]
        # Gxdown = Gx[:, 2:self.fine_height, 1:self.fine_width - 1]
        Gxleft = Gx[:, 1:self.fine_height - 1, 0:self.fine_width - 2]
        Gxright = Gx[:, 1:self.fine_height - 1, 2:self.fine_width]

        Gycenter = Gy[:, 1:self.fine_height - 1, 1:self.fine_width - 1]
        Gyup = Gy[:, 0:self.fine_height - 2, 1:self.fine_width - 1]
        Gydown = Gy[:, 2:self.fine_height, 1:self.fine_width - 1]
        # Gyleft = Gy[:, 1:self.fine_height - 1, 0:self.fine_width - 2]
        # Gyright = Gy[:, 1:self.fine_height - 1, 2:self.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))
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()
import torch
from PIL import Image

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.detach().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))
def train_GMM(model, train_loader, board,lr = 1e-4, keep_step = 100000, decay_step = 100000, save_count=500, display_count = 100, checkpoint_dir="/content/checkpoint", name = 'GMM'):
    model.cuda()
    model.train()

    L1loss = nn.L1Loss()
    Gicloss = GicLoss()

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, betas=(0.5, 0.999))

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer,
                                                  lr_lambda= lambda step: 1.0 - (max(0, step - keep_step)/float(decay_step + 1)))
    
    for step in range(keep_step + decay_step):
        iter_start_time = time.time()
        inputs = train_loader.next_batch()

        im = inputs['image'].cuda() #full image of person
        im_pose = inputs['pose_image'].cuda() # pose channels
        im_h = inputs['head'].cuda()     # person head Image 
        shape = inputs['shape'].cuda()    # blurred binary mask of person and 
        agnostic = inputs['agnostic'].cuda()  # person Representation for GMM
        c = inputs['cloth'].cuda()        # in shop cloths
        cm = inputs['cloth_mask'].cuda()    # in shop cloth mask
        im_c = inputs['parse_cloth'].cuda() # GT for GMM
        im_g = inputs['grid_image'].cuda()  # grid image for Viz.
        pcm = inputs['parse_cloth_mask'].cuda()

        grid, theta = model(agnostic, cm)
        warped_cloth = F.grid_sample(c, grid, padding_mode='border', align_corners=True)
        warped_mask = F.grid_sample(cm, grid, padding_mode='zeros', align_corners=True)
        warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros', align_corners=True)

        visuals = [[im_h, shape, im_pose],
                   [c, warped_cloth, im_c],
                   [warped_grid, (warped_cloth+im)*0.5, im]]
        
        Lwarp = L1loss(warped_mask, pcm)

        Lgic = Gicloss(grid)

        Lgic = Lgic / (grid.shape[0] * grid.shape[1] * grid.shape[2])

        loss = Lwarp + 40 * Lgic

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (step+1) % 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) % save_count == 0:
            save_checkpoint(model, os.path.join(
                checkpoint_dir, name, 'step_%06d.pth' % (step+1)))
def dir(path, name):
    name = osp.join(path, name)
    if not osp.exists(name):
        os.makedirs(name)
    return name
def test_GMM(model, test_loader, checkpoint_path = '/content/drive/Shareddrives/Virtuon/Working GMM/', name = "GMM", model_name = "PreTrainedGMM", result_dir = "results", mode = 'train'):

    model_path = osp.join(checkpoint_path, name, model_name + ".pth")
    load_checkpoint(model, model_path)

    model.cuda()
    model.eval()

    save_dir = osp.join(result_dir, name, mode)

    if not osp.exists(save_dir):
        os.makedirs(save_dir)

    warp_cloth_dir = dir(save_dir, 'warp_cloth')

    warp_mask_dir = dir(save_dir, 'warp_mask')

    result_dir = dir(save_dir, 'result_dir')

    overlayed_TPS_dir = dir(save_dir, 'overlayed_TPS')

    warped_grid_dir = dir(save_dir, 'warped_grid')

    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', align_corners=True)
        warped_mask = F.grid_sample(cm, grid, padding_mode='zeros', align_corners=True)
        warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros', align_corners=True)
        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_dir)
        save_images(warped_grid, im_names, warped_grid_dir)
        save_images(overlay, im_names, overlayed_TPS_dir)

        if (step+1) % 100 == 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)
from torch.utils.tensorboard.writer import SummaryWriter

tensorboard_dir = '/content/tensorboard'
name = 'GMM_2'
model = GMM()
train_dataset = CPDataset()
train_loader = CPDataLoader(train_dataset, batch=16)

# for viz.
if not os.path.exists(tensorboard_dir):
        os.makedirs(os.path.join(tensorboard_dir, name))
board = SummaryWriter(log_dir=os.path.join(tensorboard_dir, name))


%load_ext tensorboard
%tensorboard --logdir /content/tensorboard/GMM_2

#train_GMM(model, train_loader,board,save_count=1000,display_count=100, checkpoint_dir = '/content/drive/Shareddrives/Virtuon/Working GMM/')
test_GMM(model, train_loader)
!zip -r result.zip /content/results/GMM/train 
!cp result.zip /content/drive/Shareddrives/Virtuon/Pytorch<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1382482290932286, 'fraction_numerical': 0.02682378218686004, 'mean_word_length': 2.3368432844180576, 'pattern_counts': {'":': 1, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '475483', 'license': 'MIT', 'n_tokens_mistral': 11327, 'n_tokens_neox': 10612, 'n_words': 2004, 'path': 'research_development/Prashant/Pytorch/v3_GMM_PYTORCH.ipynb', 'repo_name': 'Pherokung/VIRTUON', 'size': '29787'}