Upload 5 files

test.py

@@ -0,0 +1,89 @@
+import time
+from options.test_options import TestOptions
+from data.data_loader_test import CreateDataLoader
+from models.networks import ResUnetGenerator, load_checkpoint
+from models.afwm import AFWM
+import torch.nn as nn
+import os
+import numpy as np
+import torch
+import cv2
+import torch.nn.functional as F
+
+opt = TestOptions().parse()
+
+start_epoch, epoch_iter = 1, 0
+
+# list human-cloth pairs
+with open('demo.txt', 'w') as file:
+    lines = [f'3.png {cloth_img_fn}\n' for cloth_img_fn in os.listdir('dataset/test_clothes')]
+    file.writelines(lines)
+
+data_loader = CreateDataLoader(opt)
+dataset = data_loader.load_data()
+dataset_size = len(data_loader)
+print(dataset_size)
+
+warp_model = AFWM(opt, 3)
+print(warp_model)
+warp_model.eval()
+warp_model.cuda()
+load_checkpoint(warp_model, opt.warp_checkpoint)
+
+gen_model = ResUnetGenerator(7, 4, 5, ngf=64, norm_layer=nn.BatchNorm2d)
+print(gen_model)
+gen_model.eval()
+gen_model.cuda()
+load_checkpoint(gen_model, opt.gen_checkpoint)
+
+total_steps = (start_epoch-1) * dataset_size + epoch_iter
+step = 0
+step_per_batch = dataset_size / opt.batchSize
+
+for epoch in range(1,2):
+
+    for i, data in enumerate(dataset, start=epoch_iter):
+        iter_start_time = time.time()
+        total_steps += opt.batchSize
+        epoch_iter += opt.batchSize
+
+        real_image = data['image']
+        clothes = data['clothes']
+        ##edge is extracted from the clothes image with the built-in function in python
+        edge = data['edge']
+        edge = torch.FloatTensor((edge.detach().numpy() > 0.5).astype(np.int))
+        clothes = clothes * edge        
+
+        flow_out = warp_model(real_image.cuda(), clothes.cuda())
+        warped_cloth, last_flow, = flow_out
+        warped_edge = F.grid_sample(edge.cuda(), last_flow.permute(0, 2, 3, 1),
+                          mode='bilinear', padding_mode='zeros')
+
+        gen_inputs = torch.cat([real_image.cuda(), warped_cloth, warped_edge], 1)
+        gen_outputs = gen_model(gen_inputs)
+        p_rendered, m_composite = torch.split(gen_outputs, [3, 1], 1)
+        p_rendered = torch.tanh(p_rendered)
+        m_composite = torch.sigmoid(m_composite)
+        m_composite = m_composite * warped_edge
+        p_tryon = warped_cloth * m_composite + p_rendered * (1 - m_composite)
+
+        path = 'results/' + opt.name
+        os.makedirs(path, exist_ok=True)
+        sub_path = path + '/PFAFN'
+        os.makedirs(sub_path,exist_ok=True)
+
+        if step % 1 == 0:
+            a = real_image.float().cuda()
+            b= clothes.cuda()
+            c = p_tryon
+            combine = torch.cat([a[0],b[0],c[0]], 2).squeeze()
+            cv_img=(combine.permute(1,2,0).detach().cpu().numpy()+1)/2
+            rgb=(cv_img*255).astype(np.uint8)
+            bgr=cv2.cvtColor(rgb,cv2.COLOR_RGB2BGR)
+            cv2.imwrite(sub_path+'/'+str(step)+'.jpg',bgr)
+
+        step += 1
+        if epoch_iter >= dataset_size:
+            break
+
+

test.sh

@@ -0,0 +1 @@
+python test.py --name demo --resize_or_crop None --batchSize 1 --gpu_ids 0

util/__init__.py

@@ -0,0 +1 @@
+# util_init

util/image_pool.py

@@ -0,0 +1,31 @@
+import random
+import torch
+from torch.autograd import Variable
+class ImagePool():
+    def __init__(self, pool_size):
+        self.pool_size = pool_size
+        if self.pool_size > 0:
+            self.num_imgs = 0
+            self.images = []
+
+    def query(self, images):
+        if self.pool_size == 0:
+            return images
+        return_images = []
+        for image in images.data:
+            image = torch.unsqueeze(image, 0)
+            if self.num_imgs < self.pool_size:
+                self.num_imgs = self.num_imgs + 1
+                self.images.append(image)
+                return_images.append(image)
+            else:
+                p = random.uniform(0, 1)
+                if p > 0.5:
+                    random_id = random.randint(0, self.pool_size-1)
+                    tmp = self.images[random_id].clone()
+                    self.images[random_id] = image
+                    return_images.append(tmp)
+                else:
+                    return_images.append(image)
+        return_images = Variable(torch.cat(return_images, 0))
+        return return_images

util/util.py

@@ -0,0 +1,94 @@
+from __future__ import print_function
+
+import torch
+from PIL import Image
+import numpy as np
+import os
+
+def tensor2im(image_tensor, imtype=np.uint8, normalize=True):
+    if isinstance(image_tensor, list):
+        image_numpy = []
+        for i in range(len(image_tensor)):
+            image_numpy.append(tensor2im(image_tensor[i], imtype, normalize))
+        return image_numpy
+    image_numpy = image_tensor.cpu().float().numpy()
+
+    image_numpy = (image_numpy + 1) / 2.0
+    image_numpy = np.clip(image_numpy, 0, 1)
+    if image_numpy.shape[2] == 1 or image_numpy.shape[2] > 3:        
+        image_numpy = image_numpy[:,:,0]
+
+    return image_numpy
+
+def tensor2label(label_tensor, n_label, imtype=np.uint8):
+    if n_label == 0:
+        return tensor2im(label_tensor, imtype)
+    label_tensor = label_tensor.cpu().float()    
+    if label_tensor.size()[0] > 1:
+        label_tensor = label_tensor.max(0, keepdim=True)[1]
+    label_tensor = Colorize(n_label)(label_tensor)
+    label_numpy = label_tensor.numpy()
+    label_numpy = label_numpy / 255.0
+
+    return label_numpy
+
+def save_image(image_numpy, image_path):
+    image_pil = Image.fromarray(image_numpy)
+    image_pil.save(image_path)
+
+def mkdirs(paths):
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            mkdir(path)
+    else:
+        mkdir(paths)
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+def uint82bin(n, count=8):
+    """returns the binary of integer n, count refers to amount of bits"""
+    return ''.join([str((n >> y) & 1) for y in range(count-1, -1, -1)])
+
+def labelcolormap(N):
+    if N == 35: # cityscape
+        cmap = np.array([(  0,  0,  0), (  0,  0,  0), (  0,  0,  0), (  0,  0,  0), (  0,  0,  0), (111, 74,  0), ( 81,  0, 81),
+                     (128, 64,128), (244, 35,232), (250,170,160), (230,150,140), ( 70, 70, 70), (102,102,156), (190,153,153),
+                     (180,165,180), (150,100,100), (150,120, 90), (153,153,153), (153,153,153), (250,170, 30), (220,220,  0),
+                     (107,142, 35), (152,251,152), ( 70,130,180), (220, 20, 60), (255,  0,  0), (  0,  0,142), (  0,  0, 70),
+                     (  0, 60,100), (  0,  0, 90), (  0,  0,110), (  0, 80,100), (  0,  0,230), (119, 11, 32), (  0,  0,142)], 
+                     dtype=np.uint8)
+    else:
+        cmap = np.zeros((N, 3), dtype=np.uint8)
+        for i in range(N):
+            r, g, b = 0, 0, 0
+            id = i
+            for j in range(7):
+                str_id = uint82bin(id)
+                r = r ^ (np.uint8(str_id[-1]) << (7-j))
+                g = g ^ (np.uint8(str_id[-2]) << (7-j))
+                b = b ^ (np.uint8(str_id[-3]) << (7-j))
+                id = id >> 3
+            cmap[i, 0] = r
+            cmap[i, 1] = g
+            cmap[i, 2] = b
+    return cmap
+
+class Colorize(object):
+    def __init__(self, n=35):
+        self.cmap = labelcolormap(n)
+        self.cmap = torch.from_numpy(self.cmap[:n])
+
+    def __call__(self, gray_image):
+        size = gray_image.size()
+        color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0)
+
+        for label in range(0, len(self.cmap)):
+            mask = (label == gray_image[0]).cpu()
+            color_image[0][mask] = self.cmap[label][0]
+            color_image[1][mask] = self.cmap[label][1]
+            color_image[2][mask] = self.cmap[label][2]
+
+        return color_image