-
Notifications
You must be signed in to change notification settings - Fork 70
/
test.py
126 lines (102 loc) · 4.37 KB
/
test.py
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
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()