def test_tom(opt, test_loader, model, board):
print('----Testing of module {} started----'.format(opt.name))
model.to(device)
model.eval()
unet_mask = UnetGenerator(25, 20, ngf=64)
load_checkpoint(unet_mask,
os.path.join(opt.checkpoint_dir, 'SEG', 'segm_final.pth'))
unet_mask.to(device)
unet_mask.eval()
gmm = GMM(opt)
load_checkpoint(gmm,
os.path.join(opt.checkpoint_dir, 'GMM', 'gmm_final.pth'))
gmm.to(device)
gmm.eval()
length = len(test_loader.data_loader)
step = 0
pbar = tqdm(total=length)
inputs = test_loader.next_batch()
while inputs is not None:
im_name = inputs['im_name']
im_h = inputs['head'].to(device)
im = inputs['image'].to(device)
agnostic = inputs['agnostic'].to(device)
c = inputs['cloth'].to(device)
# c_warp = inputs['cloth_warp'].to(device)
im_c = inputs['parse_cloth'].to(device)
im_c_mask = inputs['parse_cloth_mask'].to(device)
im_ttp = inputs['texture_t_prior'].to(device)
with torch.no_grad():
output_segm = unet_mask(torch.cat([agnostic, c], 1))
grid_zero, theta, grid_one, delta_theta = gmm(agnostic, c)
c_warp = F.grid_sample(c, grid_one, padding_mode='border')
output_segm = F.log_softmax(output_segm, dim=1)
output_argm = torch.max(output_segm, dim=1, keepdim=True)[1]
final_segm = torch.zeros(output_segm.shape).to(device).scatter(
1, output_argm, 1.0)
input_tom = torch.cat([final_segm, c_warp, im_ttp], 1)
with torch.no_grad():
output_tom = model(input_tom)
person_r = torch.tanh(output_tom[:, :3, :, :])
mask_c = torch.sigmoid(output_tom[:, 3:, :, :])
mask_c = (mask_c >= 0.5).type(torch.float)
img_tryon = mask_c * c_warp + (1 - mask_c) * person_r
visuals = [[im, c, img_tryon], [im_c, c_warp, person_r],
[im_c_mask, mask_c, im_h]]
board_add_images(board, 'combine', visuals, step + 1)
save_images(img_tryon, im_name,
osp.join(opt.dataroot, opt.datamode, 'final-output'))
inputs = test_loader.next_batch()
step += 1
pbar.update(1)
def test_tom(opt, model, inputs):
model.cuda()
model.eval()
iter_start_time = time.time()
im_pose = inputs["pose_image"]
im_h = inputs["head"]
shape = inputs["shape"]
agnostic = inputs["agnostic"].unsqueeze(0).cuda()
c = inputs["cloth"].unsqueeze(0).cuda()
cm = inputs["cloth_mask"].unsqueeze(0).cuda()
# outputs = model(torch.cat([agnostic, c], 1)) # CP-VTON
outputs = model(torch.cat([agnostic, c, cm], 1)) # CP-VTON+
p_rendered, m_composite = torch.split(outputs, 3, 1)
p_rendered = F.tanh(p_rendered)
m_composite = F.sigmoid(m_composite)
p_tryon = c * m_composite + p_rendered * (1 - m_composite)
# visuals = [[im_h, shape, im_pose],
# [c, 2*cm-1, m_composite],
# [p_rendered, p_tryon, im]]
output_dir = os.path.join(opt.result_dir, "tom")
save_images(p_tryon, ["tryon.jpg"], output_dir)
save_images(im_h, ["image_head.jpg"], output_dir)
save_images(shape, ["imshape.jpg"], output_dir)
# save_images(im_pose, ["impose.jpg"], output_dir)
save_images(m_composite, ["composite.jpg"], output_dir)
save_images(p_rendered, ["rendered.jpg"], output_dir) # For test data
t = time.time() - iter_start_time
print("time: %.3f" % (t, ), flush=True)
def test_tom(opt, test_loader, model, board):
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
save_dir = os.path.join(opt.result_dir, base_name, opt.datamode)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
try_on_dir = os.path.join(save_dir, 'try-on')
if not os.path.exists(try_on_dir):
os.makedirs(try_on_dir)
print('Dataset size: %05d!' % (len(test_loader.dataset)), flush=True)
for step, inputs in enumerate(test_loader.data_loader):
iter_start_time = time.time()
im_names = inputs['im_name']
im = inputs['image'].cuda()
im_pose = inputs['pose_image']
im_h = inputs['head']
shape = inputs['shape']
agnostic = inputs['agnostic'].cuda()
c = inputs['cloth'].cuda()
cm = inputs['cloth_mask'].cuda()
outputs = model(torch.cat([agnostic, c], 1))
p_rendered, m_composite = torch.split(outputs, 3, 1)
p_rendered = F.tanh(p_rendered)
m_composite = F.sigmoid(m_composite)
p_tryon = c * m_composite + p_rendered * (1 - m_composite)
save_images(p_tryon, im_names, try_on_dir)
'''
def test_mask_gen(opt, test_loader, model):
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
save_dir = os.path.join('./mask_gen/', base_name, opt.datamode)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
try_on_dir = os.path.join(save_dir, 'try-on')
if not os.path.exists(try_on_dir):
os.makedirs(try_on_dir)
for step, inputs in enumerate(test_loader.data_loader):
iter_start_time = time.time()
inputs = test_loader.next_batch()
im_names = inputs['im_name']
c = inputs['cloth'].cuda()
mesh = inputs['mesh'].cuda()
pose_map = inputs['pose_map'].cuda()
person_parse = inputs['person_parse'].cuda()
outputs = model(torch.cat([mesh, pose_map, c], 1))
m_composite = F.sigmoid(outputs)
save_images(m_composite, im_names, try_on_dir)
if (step + 1) % opt.display_count == 0:
t = time.time() - iter_start_time
print('step: %8d, time: %.3f' % (step + 1, t), flush=True)
print('Finished in ' + str(time.time() - start_time))
def test_step(self, batch, batch_idx):
batch = maybe_combine_frames_and_channels(self.hparams, batch)
dataset_names = batch["dataset_name"]
im_names = batch["image_name"]
if self.hparams.n_frames_total > 1:
dataset_names = get_last_item_per_batch(dataset_names)
im_names = get_last_item_per_batch(im_names)
task = "tryon" if self.hparams.tryon_list else "reconstruction"
try_on_dirs = [
osp.join(self.test_results_dir, dname, task)
for dname in dataset_names
]
# if we already did a forward-pass on this batch, skip it
save_paths = get_save_paths(try_on_dirs, im_names)
if all(osp.exists(s) for s in save_paths):
progress_bar = {"file": f"Skipping {im_names[0]}"}
else:
progress_bar = {"file": f"{im_names[0]}"}
person_inputs = get_and_cat_inputs(batch,
self.hparams.person_inputs)
cloth_inputs = get_and_cat_inputs(batch, self.hparams.cloth_inputs)
_, _, self.p_tryon, _ = self.forward(person_inputs, cloth_inputs)
# TODO CLEANUP: we get the last frame here by picking the last RGB channels;
# this is different from how it's done in training_step, which uses
# chunking and -1 indexing. We should choose one method for consistency.
save_images(self.p_tryon[:, -TryonDataset.RGB_CHANNELS:, :, ],
im_names, try_on_dirs)
result = {"progress_bar": progress_bar}
return result
def test_gmm(opt, test_loader, model, board):
if opt.use_gpu:
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
save_dir = os.path.join(opt.result_dir, base_name, opt.datamode)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
warp_cloth_dir = os.path.join(save_dir, 'warp-cloth')
if not os.path.exists(warp_cloth_dir):
os.makedirs(warp_cloth_dir)
warp_mask_dir = os.path.join(save_dir, 'warp-mask')
if not os.path.exists(warp_mask_dir):
os.makedirs(warp_mask_dir)
for step, inputs in enumerate(test_loader.data_loader):
iter_start_time = time.time()
c_names = inputs['c_name']
im = inputs['image']
im_pose = inputs['pose_image']
im_h = inputs['head']
shape = inputs['shape']
agnostic = inputs['agnostic']
c = inputs['cloth']
cm = inputs['cloth_mask']
im_c = inputs['parse_cloth']
im_g = inputs['grid_image']
if opt.use_gpu:
im = im.cuda()
im_pose = im_pose.cuda()
im_h = im_h.cuda()
shape = shape.cuda()
agnostic = agnostic.cuda()
c = c.cuda()
cm = cm.cuda()
im_c = im_c.cuda()
im_g = im_g.cuda()
grid, theta = model(agnostic, c)
warped_cloth = F.grid_sample(c, grid, padding_mode='border')
warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
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)
if (step+1) % opt.display_count == 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)
def test_tom(opt, test_loader, model, board):
if opt.cuda:
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
save_dir = os.path.join(opt.result_dir, base_name, opt.datamode)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
try_on_dir = os.path.join(save_dir, 'try-on')
if not os.path.exists(try_on_dir):
os.makedirs(try_on_dir)
print('Dataset size: %05d!' % (len(test_loader.dataset)), flush=True)
for step, inputs in enumerate(test_loader.data_loader):
iter_start_time = time.time()
im_names = inputs['im_name']
if opt.cuda:
im = inputs['image'].cuda()
else:
im = inputs['image']
im_pose = inputs['pose_image']
im_h = inputs['head']
shape = inputs['shape']
if opt.cuda:
agnostic = inputs['agnostic'].cuda()
c = inputs['cloth'].cuda()
cm = inputs['cloth_mask'].cuda()
else:
agnostic = inputs['agnostic']
c = inputs['cloth']
cm = inputs['cloth_mask']
outputs = model(torch.cat([agnostic, c],1))
p_rendered, m_composite = torch.split(outputs, 3,1)
p_rendered = F.tanh(p_rendered)
m_composite = F.sigmoid(m_composite)
p_tryon = c * m_composite + p_rendered * (1 - m_composite)
visuals = [ [im_h, shape, im_pose],
[c, 2*cm-1, m_composite],
[p_rendered, p_tryon, im]]
save_images(p_tryon, im_names, try_on_dir)
if (step+1) % opt.display_count == 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)
def test_gmm(opt, test_loader, model, board):
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
save_root = os.path.join(opt.result_dir, base_name, opt.datamode)
pbar = tqdm(enumerate(test_loader.data_loader), total=len(test_loader.data_loader))
for step, inputs in pbar:
dataset_names = inputs["dataset_name"]
# produce subfolders for each subdataset
warp_cloth_dirs = [osp.join(save_root, dname, "warp-cloth") for dname in dataset_names]
warp_mask_dirs = [osp.join(save_root, dname, "warp-mask") for dname in dataset_names]
c_names = inputs["c_name"]
# if we already did a forward-pass on this batch, skip it
save_paths = get_save_paths(c_names, warp_cloth_dirs)
if all(os.path.exists(s) for s in save_paths):
pbar.set_description(f"Skipping {c_names[0]}")
continue
pbar.set_description(c_names[0])
# unpack the rest of the data
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()
# forward pass
grid, theta = model(agnostic, c)
warped_cloth = F.grid_sample(c, grid, padding_mode='border')
warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
# save images
save_images(warped_cloth, c_names, warp_cloth_dirs)
save_images(warped_mask*2-1, c_names, warp_mask_dirs)
if opt.tensorboard_dir and (step+1) % opt.display_count == 0:
visuals = [[im_h, shape, im_pose],
[c, warped_cloth, im_c],
[warped_grid, (warped_cloth + im) * 0.5, im]]
board_add_images(board, 'combine', visuals, step+1)
def test_gmm(opt, test_loader, model):
#model.cuda()
model.eval()
save_dir = opt.result_dir
if not os.path.exists(save_dir):
os.makedirs(save_dir)
warp_cloth_dir = os.path.join(save_dir, 'warp-cloth')
if not os.path.exists(warp_cloth_dir):
os.makedirs(warp_cloth_dir)
warp_mask_dir = os.path.join(save_dir, 'warp-mask')
if not os.path.exists(warp_mask_dir):
os.makedirs(warp_mask_dir)
for step, inputs in enumerate(test_loader.data_loader):
iter_start_time = time.time()
c_names = inputs['c_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()
im = inputs['image']
im_pose = inputs['pose_image']
im_h = inputs['head']
shape = inputs['shape']
agnostic = inputs['agnostic']
c = inputs['cloth']
cm = inputs['cloth_mask']
im_c = inputs['parse_cloth']
im_g = inputs['grid_image']
grid, theta = model(agnostic, c)
warped_cloth = F.grid_sample(c, grid, padding_mode='border')
warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
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)
def test_tom(opt, test_loader, model, board):
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
save_root = os.path.join(opt.result_dir, base_name, opt.datamode)
print('Dataset size: %05d!' % (len(test_loader.dataset)), flush=True)
pbar = tqdm(enumerate(test_loader.data_loader), total=len(test_loader.data_loader))
for step, inputs in pbar:
dataset_names = inputs["dataset_name"]
# use subfolders for each subdataset
try_on_dirs = [osp.join(save_root, dname, "try-on") for dname in dataset_names]
im_names = inputs['im_name']
# if we already did a forward-pass on this batch, skip it
save_paths = get_save_paths(im_names, try_on_dirs)
if all(os.path.exists(s) for s in save_paths):
tqdm.write(f"Skipping {save_paths}")
continue
pbar.set_description(im_names[0])
im = inputs['image'].cuda()
im_pose = inputs['pose_image']
im_h = inputs['head']
shape = inputs['shape']
agnostic = inputs['agnostic'].cuda()
c = inputs['cloth'].cuda()
cm = inputs['cloth_mask'].cuda()
outputs = model(torch.cat([agnostic, c],1))
p_rendered, m_composite = torch.split(outputs, 3,1)
p_rendered = F.tanh(p_rendered)
m_composite = F.sigmoid(m_composite)
p_tryon = c * m_composite + p_rendered * (1 - m_composite)
visuals = [ [im_h, shape, im_pose],
[c, 2*cm-1, m_composite],
[p_rendered, p_tryon, im]]
save_images(p_tryon, im_names, try_on_dirs)
if opt.tensorboard_dir and (step+1) % opt.display_count == 0:
board_add_images(board, 'combine', visuals, step+1)
def test_gmm(opt, test_loader, model, board):
# load model
model.cuda()
model.eval()
# make dirs
base_name = os.path.basename(opt.checkpoint)
warp_cloth_dir = os.path.join(opt.result_dir, opt.name, base_name,
opt.data_list.split('.')[0], 'warp-cloth')
os.makedirs(warp_cloth_dir, exist_ok=True)
warp_mask_dir = os.path.join(opt.result_dir, opt.name, base_name,
opt.data_list.split('.')[0], 'warp-mask')
os.makedirs(warp_mask_dir, exist_ok=True)
# test loop
for step, inputs in enumerate(iter(test_loader)):
iter_start_time = time.time()
c_names = inputs['c_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()
grid, theta = model(agnostic, c)
warped_cloth = F.grid_sample(c, grid, padding_mode='border')
warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
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)
if (step + 1) % opt.display_count == 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)
def test_gmm(opt, test_loader, model, board):
print('----Testing of module {} started----'.format(opt.name))
model.to(device)
model.eval()
length = len(test_loader.data_loader)
step = 0
pbar = tqdm(total=length)
inputs = test_loader.next_batch()
while inputs is not None:
im_name = inputs['im_name']
im = inputs['image'].to(device)
im_h = inputs['head'].to(device)
agnostic = inputs['agnostic'].to(device)
c = inputs['cloth'].to(device)
im_c = inputs['parse_cloth'].to(device)
im_g = inputs['grid_image'].to(device)
with torch.no_grad():
grid_zero, theta, grid_one, delta_theta = model(agnostic, c)
warped_coarse_cloth = F.grid_sample(c,
grid_zero,
padding_mode='border')
warped_cloth = F.grid_sample(c, grid_one, padding_mode='border')
warped_coarse_grid = F.grid_sample(im_g,
grid_zero,
padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid_one, padding_mode='zeros')
visuals = [[im_h, im, (warped_cloth + im) * 0.5],
[warped_grid, warped_coarse_grid, warped_grid],
[im_c, warped_coarse_cloth, warped_cloth]]
board_add_images(board, 'combine', visuals, step + 1)
save_images(warped_cloth, im_name,
osp.join(opt.dataroot, opt.datamode, 'wrap-cloth'))
inputs = test_loader.next_batch()
step += 1
pbar.update(1)
def test_step(self, batch, batch_idx):
batch = maybe_combine_frames_and_channels(self.hparams, batch)
dataset_names = batch["dataset_name"]
# produce subfolders for each subdataset
warp_cloth_dirs = [
osp.join(self.test_results_dir, dname, "warp-cloth")
for dname in dataset_names
]
warp_mask_dirs = [
osp.join(self.test_results_dir, dname, "warp-mask")
for dname in dataset_names
]
c_names = batch["cloth_name"]
# if we already did a forward-pass on this batch, skip it
save_paths = get_save_paths(warp_cloth_dirs, c_names)
if all(osp.exists(s) for s in save_paths):
progress_bar = {"file": f"Skipping {c_names[0]}"}
else:
progress_bar = {"file": c_names[0]}
# unpack the the data
c = batch["cloth"]
cm = batch["cloth_mask"]
im_g = batch["grid_vis"]
person_inputs = get_and_cat_inputs(batch,
self.hparams.person_inputs)
cloth_inputs = get_and_cat_inputs(batch, self.hparams.cloth_inputs)
# forward pass
grid, theta = self.forward(person_inputs, cloth_inputs)
self.warped_cloth = F.grid_sample(c, grid, padding_mode="border")
warped_mask = F.grid_sample(cm, grid, padding_mode="zeros")
self.warped_grid = F.grid_sample(im_g, grid, padding_mode="zeros")
# save images
save_images(self.warped_cloth, c_names, warp_cloth_dirs)
save_images(warped_mask * 2 - 1, c_names, warp_mask_dirs)
result = {"progress_bar": progress_bar}
return result
def test_hpm(opt, test_loader, model, board):
model.cuda()
model.eval()
image_seg_dir = os.path.join(save_dir, 'image-seg')
if not os.path.exists(image_seg_dir):
os.makedirs(image_seg_dir)
warp_mask_dir = os.path.join(save_dir, 'warp-mask')
if not os.path.exists(image_seg_dir):
os.makedirs(image_seg_dir)
for step, inputs in enumerate(test_loader.data_loader):
agnostic = inputs['agnostic'].cuda()
c = inputs['cloth'].cuda()
#generate segmentation
segmentation = model(torch.cat([agnostic, c], 1))
save_images(segmentation, im_names, image_seg_dir)
#generate mask
warped_mask = segmentation
save_images(warped_mask * 2 - 1, c_names, warp_mask_dir)
def test_tom(opt, test_loader, model, board):
# load model
model.cuda()
model.eval()
# make dirs
base_name = os.path.basename(opt.checkpoint)
try_on_dir = os.path.join(opt.result_dir, opt.name, base_name, 'try_on')
os.makedirs(try_on_dir, exist_ok=True)
# test loop
for step, inputs in enumerate(iter(test_loader)):
iter_start_time = time.time()
im_names = inputs['im_name']
im = inputs['image'].cuda()
im_pose = inputs['pose_image']
im_h = inputs['head']
shape = inputs['shape']
agnostic = inputs['agnostic'].cuda()
c = inputs['cloth'].cuda()
cm = inputs['cloth_mask'].cuda()
outputs = model(torch.cat([agnostic, c], 1))
p_rendered, m_composite = torch.split(outputs, 3, 1)
p_rendered = torch.tanh(p_rendered)
m_composite = torch.sigmoid(m_composite)
p_tryon = c * m_composite + p_rendered * (1 - m_composite)
visuals = [[im_h, shape, im_pose], [c, 2 * cm - 1, m_composite],
[p_rendered, p_tryon, im]]
save_images(p_tryon, im_names, try_on_dir)
if (step + 1) % opt.display_count == 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)
def test_gmm(opt, test_loader, model, board):
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
save_dir = os.path.join(opt.result_dir, base_name, opt.datamode)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
warp_cloth_dir = os.path.join(save_dir, 'warp-cloth')
if not os.path.exists(warp_cloth_dir):
os.makedirs(warp_cloth_dir)
for step, inputs in enumerate(test_loader.data_loader):
iter_start_time = time.time()
c_names = inputs['c_name']
wm = inputs['warped_mask'].cuda()
c = inputs['cloth'].cuda()
grid, theta = model(wm, c)
warped_cloth = F.grid_sample(c, grid, padding_mode='border')
save_images(warped_cloth, c_names, warp_cloth_dir)
def test_gmm(opt, test_loader, model, board):
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
save_dir = os.path.join(opt.result_dir, base_name, opt.datamode)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
warp_cloth_dir = os.path.join(save_dir, 'warp-cloth')
if not os.path.exists(warp_cloth_dir):
os.makedirs(warp_cloth_dir)
warp_mask_dir = os.path.join(save_dir, 'warp-mask')
if not os.path.exists(warp_mask_dir):
os.makedirs(warp_mask_dir)
for step, inputs in enumerate(test_loader.data_loader):
iter_start_time = time.time()
c_names = inputs['c_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()
grid, theta = model(agnostic, c)
warped_cloth = F.grid_sample(c, grid, padding_mode='border')
warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
save_images(warped_cloth, c_names, warp_cloth_dir)
save_images(warped_mask * 2 - 1, c_names, warp_mask_dir)
'''
def predict_gmm(opt, model, inputs, save_intermediate=False):
model.cuda()
model.eval()
person_image = inputs["image"].cuda()
cloth_mask = inputs["cloth_mask"].unsqueeze(0).cuda()
cloth_orig = inputs["cloth"].unsqueeze(0).cuda()
image_grid = inputs["grid_image"].unsqueeze(0).cuda()
agnostic = inputs["agnostic"].unsqueeze(0).cuda()
grid, theta = model(agnostic, cloth_mask)
warped_cloth = F.grid_sample(cloth_orig, grid, padding_mode="border")
warped_mask = F.grid_sample(cloth_mask, grid, padding_mode="zeros")
warped_grid = F.grid_sample(image_grid, grid, padding_mode="zeros")
overlay = 0.7 * warped_cloth + 0.3 * person_image
if save_intermediate:
save_images(warped_cloth, ["warped_cloth.jpg"], "output/debug")
save_images(warped_mask, ["warped_mask.jpg"], "output/debug")
save_images(warped_grid, ["warped_grid.jpg"], "output/debug")
save_images(overlay, ["overlay.jpg"], "output/debug")
return overlay
def test_gmm(opt, test_loader, model, board):
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
name = opt.name
save_dir = os.path.join(opt.result_dir, name, opt.datamode)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
warp_cloth_dir = os.path.join(save_dir, 'warp-cloth')
if not os.path.exists(warp_cloth_dir):
os.makedirs(warp_cloth_dir)
warp_mask_dir = os.path.join(save_dir, 'warp-mask')
if not os.path.exists(warp_mask_dir):
os.makedirs(warp_mask_dir)
warp_grid_dir = os.path.join(save_dir, 'warp-grid')
if not os.path.exists(warp_grid_dir):
os.makedirs(warp_grid_dir)
Ipersp_dir = os.path.join(save_dir, 'warp-cloth_affine')
if not os.path.exists(Ipersp_dir):
os.makedirs(Ipersp_dir)
warped_mask_affine_dir = os.path.join(save_dir, 'warped_mask_affine')
if not os.path.exists(warped_mask_affine_dir):
os.makedirs(warped_mask_affine_dir)
warped_grid_affine_dir = os.path.join(save_dir, 'warped_grid_affine')
if not os.path.exists(warped_grid_affine_dir):
os.makedirs(warped_grid_affine_dir)
overlayed_affine_dir = os.path.join(save_dir, 'overlayed_affine')
if not os.path.exists(overlayed_affine_dir):
os.makedirs(overlayed_affine_dir)
overlayed_TPS_dir = os.path.join(save_dir, 'overlayed_TPS')
if not os.path.exists(overlayed_TPS_dir):
os.makedirs(overlayed_TPS_dir)
for step, inputs in enumerate(test_loader.data_loader):
iter_start_time = time.time()
im_names = inputs['im_name']
c_names = inputs['c_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()
gridtps, thetatps, Ipersp, gridaffine, thetaaffine = model(agnostic, c)
warped_grid_affine = F.grid_sample(im_g,
gridaffine,
padding_mode='zeros')
warped_mask_affine = F.grid_sample(cm,
gridaffine,
padding_mode='zeros')
warped_mask = F.grid_sample(warped_mask_affine,
gridtps,
padding_mode='zeros')
warped_cloth = F.grid_sample(Ipersp, gridtps, padding_mode='border') #
warped_grid = F.grid_sample(warped_grid_affine,
gridtps,
padding_mode='zeros')
overlayed_affine = (Ipersp * 0.7 + im * 0.3)
overlayed_TPS = (warped_cloth * 0.7 + im * 0.3)
visuals = [[im_h, shape, im_pose], [c, warped_cloth, im_c],
[warped_grid, (warped_cloth * 0.7 + im * 0.3), im],
[Ipersp, (Ipersp * 0.7 + im * 0.3), warped_grid_affine]]
save_images(warped_cloth, im_names, warp_cloth_dir)
save_images(warped_mask * 2 - 1, im_names, warp_mask_dir)
save_images(warped_grid, im_names, warp_grid_dir)
save_images(Ipersp, im_names, Ipersp_dir)
save_images(warped_mask_affine * 2 - 1, im_names,
warped_mask_affine_dir)
save_images(warped_grid_affine, im_names, warped_grid_affine_dir)
save_images(overlayed_affine, im_names, overlayed_affine_dir)
save_images(overlayed_TPS, im_names, overlayed_TPS_dir)
if (step + 1) % opt.display_count == 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)
def test_gmm(opt, test_loader, model, board):
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
name = opt.name
save_dir = os.path.join(opt.result_dir, name, opt.datamode)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
warp_cloth_dir = os.path.join(save_dir, 'warp-cloth')
if not os.path.exists(warp_cloth_dir):
os.makedirs(warp_cloth_dir)
warp_mask_dir = os.path.join(save_dir, 'warp-mask')
if not os.path.exists(warp_mask_dir):
os.makedirs(warp_mask_dir)
result_dir1 = os.path.join(save_dir, 'result_dir')
if not os.path.exists(result_dir1):
os.makedirs(result_dir1)
overlayed_TPS_dir = os.path.join(save_dir, 'overlayed_TPS')
if not os.path.exists(overlayed_TPS_dir):
os.makedirs(overlayed_TPS_dir)
warped_grid_dir = os.path.join(save_dir, 'warped_grid')
if not os.path.exists(warped_grid_dir):
os.makedirs(warped_grid_dir)
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')
warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
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_dir1)
save_images(warped_grid, im_names, warped_grid_dir)
save_images(overlay, im_names, overlayed_TPS_dir)
if (step + 1) % opt.display_count == 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)
def test_gmm(opt, test_loader, model, board):
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
save_dir = os.path.join(opt.result_dir, base_name, opt.datamode)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
warp_cloth_dir = os.path.join(save_dir, 'warp-cloth')
if not os.path.exists(warp_cloth_dir):
os.makedirs(warp_cloth_dir)
warp_mask_dir = os.path.join(save_dir, 'warp-mask')
if not os.path.exists(warp_mask_dir):
os.makedirs(warp_mask_dir)
try:
if (opt.datamode == "train"):
f = open('data/tom_train_pairs.txt', 'w')
elif (opt.datamode == "test"):
f = open('data/tom_test_pairs.txt', 'w')
f.write("")
f.close()
except:
pass
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()
grid, theta = model(agnostic, c)
warped_cloth = F.grid_sample(c, grid, padding_mode='border')
warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
visuals = [[im_h, shape, im_pose], [c, warped_cloth, im_c],
[warped_grid, (warped_cloth + im) * 0.5, im]]
if (opt.datamode == "train"):
f = open('data/tom_train_pairs.txt', 'a')
elif (opt.datamode == "test"):
f = open('data/tom_test_pairs.txt', 'a')
i = 0
for im_name in im_names:
new_name = im_name[:-4] + "_00.jpg"
f.write(im_name + " " + new_name + "\n")
im_names[i] = new_name
i = i + 1
save_images(warped_cloth, im_names, warp_cloth_dir)
save_images(warped_mask * 2 - 1, im_names, warp_mask_dir)
if (step + 1) % opt.display_count == 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)
def test_no_background_refined_gmm(opt, test_loader, model, board):
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
save_dir = os.path.join(opt.result_dir, base_name, opt.datamode)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
warp_cloth_dir = os.path.join(save_dir, 'warp-cloth')
if not os.path.exists(warp_cloth_dir):
os.makedirs(warp_cloth_dir)
warp_mask_dir = os.path.join(save_dir, 'warp-mask')
if not os.path.exists(warp_mask_dir):
os.makedirs(warp_mask_dir)
warp_refined_dir = os.path.join(save_dir, 'warp-refined')
if not os.path.exists(warp_refined_dir):
os.makedirs(warp_refined_dir)
no_background_c_dir = os.path.join(save_dir, 'no_background_c')
if not os.path.exists(no_background_c_dir):
os.makedirs(no_background_c_dir)
for step, inputs in enumerate(test_loader.data_loader):
iter_start_time = time.time()
c_names = inputs['c_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()
no_background_c = inputs['no_background_cloth'].cuda()
cm = inputs['cloth_mask'].cuda()
im_c = inputs['parse_cloth'].cuda()
im_g = inputs['grid_image'].cuda()
c_point_plane = inputs['cloth_points'].cuda()
p_point_plane = inputs['person_points'].cuda()
grid, theta, warped_cloth, outputs = model(agnostic, no_background_c)
#warped_cloth = F.grid_sample(c, grid, padding_mode='border')
warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
#compute_c_point_plane = compute_grid_point(p_point_plane, grid)
c_rendered, m_composite = torch.split(outputs, 3,1)
c_rendered = F.tanh(c_rendered)
m_composite = F.sigmoid(m_composite)
c_result = warped_cloth * m_composite + c_rendered * (1 - m_composite)
visuals = [ [im_h, shape, im_pose],
[no_background_c, warped_cloth, im_c],
[warped_grid, (warped_cloth+im)*0.5, im],
[m_composite, (c_result+im)*0.5, c_result]]
save_images(warped_cloth, c_names, warp_cloth_dir)
save_images(warped_mask*2-1, c_names, warp_mask_dir)
#print("warp_refined_dir = ", warp_refined_dir)
#assert os.path.exists(warp_refined_dir)
save_images(no_background_c, c_names, no_background_c_dir)
save_images(c_result, c_names, warp_refined_dir)
if (step+1) % opt.display_count == 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)
def test_tom(opt, test_loader, model, board):
model.cuda()
model.eval()
base_name = os.path.basename(opt.checkpoint)
save_dir = os.path.join(opt.result_dir, base_name, opt.datamode)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
try_on_dir = os.path.join(save_dir, 'try-on')
if not os.path.exists(try_on_dir):
os.makedirs(try_on_dir)
render_dir = os.path.join(save_dir, 'render_dir')
if not os.path.exists(render_dir):
os.makedirs(render_dir)
alpha_top_dir = os.path.join(save_dir, 'alpha_top_dir')
if not os.path.exists(alpha_top_dir):
os.makedirs(alpha_top_dir)
alpha_bottom_dir = os.path.join(save_dir, 'alpha_bottom_dir')
if not os.path.exists(alpha_bottom_dir):
os.makedirs(alpha_bottom_dir)
print('Dataset size: %05d!' % (len(test_loader.dataset)), flush=True)
for step, inputs in enumerate(test_loader.data_loader):
iter_start_time = time.time()
im_names = inputs['im_name']
im = inputs['image'].cuda()
im_pose = inputs['pose_image']
im_h = inputs['head']
shape = inputs['shape']
agnostic = inputs['agnostic'].cuda()
c = inputs['cloth'].cuda()
cm = inputs['cloth_mask'].cuda()
cbottom = inputs['cloth_bottom'].cuda()
cmbottom = inputs['cloth_mask_bottom'].cuda()
outputs = model(torch.cat([agnostic, c, cbottom], 1))
p_rendered, m_composite = torch.split(outputs, 3, 1)
m_composite, m_composite_bottom = torch.split(m_composite, 1, 1)
p_rendered = F.tanh(p_rendered)
m_composite = F.sigmoid(m_composite)
m_composite_bottom = F.sigmoid(m_composite_bottom)
p_tryon = c * m_composite + p_rendered * (
1 - m_composite -
m_composite_bottom) + cbottom * m_composite_bottom
#print("m_composite:", m_composite.shape, ", max: ", m_composite.max, ", min: ", m_composite.min)
visuals = [[im_h, shape, im_pose],
[c, cm * 2 - 1, m_composite * 2 - 1],
[cbottom, cmbottom * 2 - 1, m_composite_bottom * 2 - 1],
[p_rendered, p_tryon, im]]
save_images(p_tryon, im_names, try_on_dir, isCloth=True)
save_images(p_rendered, im_names, render_dir, isCloth=True)
save_images(m_composite * 2 - 1, im_names, alpha_top_dir, isCloth=True)
save_images(m_composite_bottom * 2 - 1,
im_names,
alpha_bottom_dir,
isCloth=True)
if (step + 1) % opt.display_count == 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)
