Spaces:
Runtime error
Runtime error
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() | |