Spaces:
Runtime error
Runtime error
File size: 4,476 Bytes
4a285f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
import time
from collections import OrderedDict
from options.test_options import TestOptions
from data.data_loader import CreateDataLoader
from models.models import create_model
import util.util as util
import os
import numpy as np
import torch
from torch.autograd import Variable
#from tensorboardX import SummaryWriter
import cv2
#writer = SummaryWriter('runs/G1G2')
SIZE = 320
NC = 14
def generate_label_plain(inputs):
size = inputs.size()
pred_batch = []
for input in inputs:
input = input.view(1, NC, 256, 192)
pred = np.squeeze(input.data.max(1)[1].cpu().numpy(), axis=0)
pred_batch.append(pred)
pred_batch = np.array(pred_batch)
pred_batch = torch.from_numpy(pred_batch)
label_batch = pred_batch.view(size[0], 1, 256, 192)
return label_batch
def generate_label_color(inputs):
label_batch = []
for i in range(len(inputs)):
label_batch.append(util.tensor2label(inputs[i], NC))
label_batch = np.array(label_batch)
label_batch = label_batch * 2 - 1
input_label = torch.from_numpy(label_batch)
return input_label
def complete_compose(img, mask, label):
label = label.cpu().numpy()
M_f = label > 0
M_f = M_f.astype(np.int)
M_f = torch.FloatTensor(M_f).cuda()
masked_img = img*(1-mask)
M_c = (1-mask.cuda())*M_f
M_c = M_c+torch.zeros(img.shape).cuda() # broadcasting
return masked_img, M_c, M_f
def compose(label, mask, color_mask, edge, color, noise):
masked_label = label*(1-mask)
masked_edge = mask*edge
masked_color_strokes = mask*(1-color_mask)*color
masked_noise = mask*noise
return masked_label, masked_edge, masked_color_strokes, masked_noise
def changearm(old_label):
label = old_label
arm1 = torch.FloatTensor((old_label.cpu().numpy() == 11).astype(np.int))
arm2 = torch.FloatTensor((old_label.cpu().numpy() == 13).astype(np.int))
noise = torch.FloatTensor((old_label.cpu().numpy() == 7).astype(np.int))
label = label*(1-arm1)+arm1*4
label = label*(1-arm2)+arm2*4
label = label*(1-noise)+noise*4
return label
def main():
os.makedirs('sample', exist_ok=True)
opt = TestOptions().parse()
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('# Inference images = %d' % dataset_size)
model = create_model(opt)
for i, data in enumerate(dataset):
# add gaussian noise channel
# wash the label
t_mask = torch.FloatTensor(
(data['label'].cpu().numpy() == 7).astype(np.float))
#
# data['label'] = data['label'] * (1 - t_mask) + t_mask * 4
mask_clothes = torch.FloatTensor(
(data['label'].cpu().numpy() == 4).astype(np.int))
mask_fore = torch.FloatTensor(
(data['label'].cpu().numpy() > 0).astype(np.int))
img_fore = data['image'] * mask_fore
img_fore_wc = img_fore * mask_fore
all_clothes_label = changearm(data['label'])
############## Forward Pass ######################
fake_image, warped_cloth, refined_cloth = model(Variable(data['label'].cuda()), Variable(data['edge'].cuda()), Variable(img_fore.cuda()), Variable(
mask_clothes.cuda()), Variable(data['color'].cuda()), Variable(all_clothes_label.cuda()), Variable(data['image'].cuda()), Variable(data['pose'].cuda()), Variable(data['image'].cuda()), Variable(mask_fore.cuda()))
# make output folders
output_dir = os.path.join(opt.results_dir, opt.phase)
fake_image_dir = os.path.join(output_dir, 'try-on')
os.makedirs(fake_image_dir, exist_ok=True)
warped_cloth_dir = os.path.join(output_dir, 'warped_cloth')
os.makedirs(warped_cloth_dir, exist_ok=True)
refined_cloth_dir = os.path.join(output_dir, 'refined_cloth')
os.makedirs(refined_cloth_dir, exist_ok=True)
# save output
for j in range(opt.batchSize):
print("Saving", data['name'][j])
util.save_tensor_as_image(fake_image[j],
os.path.join(fake_image_dir, data['name'][j]))
util.save_tensor_as_image(warped_cloth[j],
os.path.join(warped_cloth_dir, data['name'][j]))
util.save_tensor_as_image(refined_cloth[j],
os.path.join(refined_cloth_dir, data['name'][j]))
if __name__ == '__main__':
main()
|