Closed MotiBaadror closed 4 years ago
Hi @MotiBaadror , looks like there's an artifact with cloth shape in neck area. Can you please share your cp_dataset.py
? Also, I think CIHP_PGN segmentation does not have labels for hand-band or wrist-watch, so it will be absent in the final results.
Hi @minar09 , Thank you for your response. here is my cp_dataset.py. If I will not get the watch in the result. It would be fine. I just want to solve try-on problem. Also, I am curious that, why the third image output is blurred near the waist for output?
# coding=utf-8
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
from PIL import Image
from PIL import ImageDraw
import os.path as osp
import numpy as np
import json
class CPDataset(data.Dataset):
"""Dataset for CP-VTON+.
"""
def __init__(self, opt):
super(CPDataset, self).__init__()
# base setting
self.opt = opt
self.root = opt.dataroot
self.datamode = opt.datamode # train or test or self-defined
self.stage = opt.stage # GMM or TOM
self.data_list = opt.data_list
self.fine_height = opt.fine_height
self.fine_width = opt.fine_width
self.radius = opt.radius
self.data_path = osp.join(opt.dataroot, opt.datamode)
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# load data list
im_names = []
c_names = []
with open(osp.join(opt.dataroot, opt.data_list), 'r') as f:
for line in f.readlines():
im_name, c_name = line.strip().split()
im_names.append(im_name)
c_names.append(c_name)
self.im_names = im_names
self.c_names = c_names
def name(self):
return "CPDataset"
def __getitem__(self, index):
c_name = self.c_names[index]
im_name = self.im_names[index]
if self.stage == 'GMM':
c = Image.open(osp.join(self.data_path, 'cloth', c_name))
cm = Image.open(osp.join(self.data_path, 'cloth-mask', c_name)).convert('L')
else:
c = Image.open(osp.join(self.data_path, 'warp-cloth', im_name)) # c_name, if that is used when saved
cm = Image.open(osp.join(self.data_path, 'warp-mask', im_name)).convert('L') # c_name, if that is used when saved
c = self.transform(c) # [-1,1]
cm_array = np.array(cm)
cm_array = (cm_array >= 128).astype(np.float32)
cm = torch.from_numpy(cm_array) # [0,1]
cm.unsqueeze_(0)
# person image
im = Image.open(osp.join(self.data_path, 'image', im_name))
im = self.transform(im) # [-1,1]
"""
LIP labels
[(0, 0, 0), # 0=Background
(128, 0, 0), # 1=Hat
(255, 0, 0), # 2=Hair
(0, 85, 0), # 3=Glove
(170, 0, 51), # 4=SunGlasses
(255, 85, 0), # 5=UpperClothes
(0, 0, 85), # 6=Dress
(0, 119, 221), # 7=Coat
(85, 85, 0), # 8=Socks
(0, 85, 85), # 9=Pants
(85, 51, 0), # 10=Jumpsuits
(52, 86, 128), # 11=Scarf
(0, 128, 0), # 12=Skirt
(0, 0, 255), # 13=Face
(51, 170, 221), # 14=LeftArm
(0, 255, 255), # 15=RightArm
(85, 255, 170), # 16=LeftLeg
(170, 255, 85), # 17=RightLeg
(255, 255, 0), # 18=LeftShoe
(255, 170, 0) # 19=RightShoe
(170, 170, 50) # 20=Skin/Neck/Chest (Newly added after running dataset_neck_skin_correction.py)
]
"""
# load parsing image
parse_name = im_name.replace('.jpg', '.png')
im_parse = Image.open(osp.join(self.data_path, 'image-parse-new', parse_name)).convert('L') # read segmentation
parse_array = np.array(im_parse) # convert to numpy array
parse_shape = (parse_array > 0).astype(np.float32) # get binary body shape
# im_parse = Image.open(
# # osp.join(self.data_path, 'image-parse', parse_name)).convert('L')
# osp.join(self.data_path, 'image-parse-new', parse_name)).convert('L') # updated new segmentation
# parse_array = np.array(im_parse)
# im_mask = Image.open(
# osp.join(self.data_path, 'image-mask', parse_name)).convert('L')
# mask_array = np.array(im_mask)
# # parse_shape = (parse_array > 0).astype(np.float32) # CP-VTON body shape
# # Get shape from body mask (CP-VTON+)
# parse_shape = (mask_array > 0).astype(np.float32)
if self.stage == 'GMM':
parse_head = (parse_array == 1).astype(np.float32) + \
(parse_array == 4).astype(np.float32) + \
(parse_array == 13).astype(
np.float32) # CP-VTON+ GMM input (reserved regions)
else:
parse_head = (parse_array == 1).astype(np.float32) + \
(parse_array == 2).astype(np.float32) + \
(parse_array == 4).astype(np.float32) + \
(parse_array == 9).astype(np.float32) + \
(parse_array == 12).astype(np.float32) + \
(parse_array == 13).astype(np.float32) + \
(parse_array == 16).astype(np.float32) + \
(parse_array == 17).astype(
np.float32) # CP-VTON+ TOM input (reserved regions)
parse_cloth = (parse_array == 5).astype(np.float32) + \
(parse_array == 6).astype(np.float32) + \
(parse_array == 7).astype(np.float32) # upper-clothes labels
# shape downsample
parse_shape_ori = Image.fromarray((parse_shape*255).astype(np.uint8))
parse_shape = parse_shape_ori.resize(
(self.fine_width//16, self.fine_height//16), Image.BILINEAR)
parse_shape = parse_shape.resize(
(self.fine_width, self.fine_height), Image.BILINEAR)
parse_shape_ori = parse_shape_ori.resize(
(self.fine_width, self.fine_height), Image.BILINEAR)
shape_ori = self.transform(parse_shape_ori) # [-1,1]
shape = self.transform(parse_shape) # [-1,1]
phead = torch.from_numpy(parse_head) # [0,1]
# phand = torch.from_numpy(parse_hand) # [0,1]
pcm = torch.from_numpy(parse_cloth) # [0,1]
# upper cloth
im_c = im * pcm + (1 - pcm) # [-1,1], fill 1 for other parts
im_h = im * phead - (1 - phead) # [-1,1], fill 0 for other parts
# load pose points
pose_name = im_name.replace('.jpg', '_keypoints.json')
with open(osp.join(self.data_path, 'pose', pose_name), 'r') as f:
pose_label = json.load(f)
pose_data = pose_label['people'][0]['pose_keypoints_2d']
pose_data = np.array(pose_data)
pose_data = pose_data.reshape((-1, 3))
point_num = pose_data.shape[0]
pose_map = torch.zeros(point_num, self.fine_height, self.fine_width)
r = self.radius
im_pose = Image.new('L', (self.fine_width, self.fine_height))
pose_draw = ImageDraw.Draw(im_pose)
for i in range(point_num):
one_map = Image.new('L', (self.fine_width, self.fine_height))
draw = ImageDraw.Draw(one_map)
pointx = pose_data[i, 0]
pointy = pose_data[i, 1]
if pointx > 1 and pointy > 1:
draw.rectangle((pointx-r, pointy-r, pointx +
r, pointy+r), 'white', 'white')
pose_draw.rectangle(
(pointx-r, pointy-r, pointx+r, pointy+r), 'white', 'white')
one_map = self.transform(one_map)
pose_map[i] = one_map[0]
# just for visualization
im_pose = self.transform(im_pose)
# cloth-agnostic representation
agnostic = torch.cat([shape, im_h, pose_map], 0)
if self.stage == 'GMM':
im_g = Image.open('grid.png')
im_g = self.transform(im_g)
else:
im_g = ''
pcm.unsqueeze_(0) # CP-VTON+
result = {
'c_name': c_name, # for visualization
'im_name': im_name, # for visualization or ground truth
'cloth': c, # for input
'cloth_mask': cm, # for input
'image': im, # for visualization
'agnostic': agnostic, # for input
'parse_cloth': im_c, # for ground truth
'shape': shape, # for visualization
'head': im_h, # for visualization
'pose_image': im_pose, # for visualization
'grid_image': im_g, # for visualization
'parse_cloth_mask': pcm, # for CP-VTON+, TOM input
'shape_ori': shape_ori, # original body shape without resize
}
return result
def __len__(self):
return len(self.im_names)
class CPDataLoader(object):
def __init__(self, opt, dataset):
super(CPDataLoader, self).__init__()
if opt.shuffle:
train_sampler = torch.utils.data.sampler.RandomSampler(dataset)
else:
train_sampler = None
self.data_loader = torch.utils.data.DataLoader(
dataset, batch_size=opt.batch_size, shuffle=(
train_sampler is None),
num_workers=opt.workers, pin_memory=True, sampler=train_sampler)
self.dataset = dataset
self.data_iter = self.data_loader.__iter__()
def next_batch(self):
try:
batch = self.data_iter.__next__()
except StopIteration:
self.data_iter = self.data_loader.__iter__()
batch = self.data_iter.__next__()
return batch
if __name__ == "__main__":
print("Check the dataset for geometric matching module!")
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--dataroot", default="data")
parser.add_argument("--datamode", default="train")
parser.add_argument("--stage", default="GMM")
parser.add_argument("--data_list", default="train_pairs.txt")
parser.add_argument("--fine_width", type=int, default=192)
parser.add_argument("--fine_height", type=int, default=256)
parser.add_argument("--radius", type=int, default=3)
parser.add_argument("--shuffle", action='store_true',
help='shuffle input data')
parser.add_argument('-b', '--batch-size', type=int, default=4)
parser.add_argument('-j', '--workers', type=int, default=1)
opt = parser.parse_args()
dataset = CPDataset(opt)
data_loader = CPDataLoader(opt, dataset)
print('Size of the dataset: %05d, dataloader: %04d'
% (len(dataset), len(data_loader.data_loader)))
first_item = dataset.__getitem__(0)
first_batch = data_loader.next_batch()
from IPython import embed
embed()
@MotiBaadror , looks like code is okay. Maybe the network is somehow still getting the source cloth shape in GMM so the first warped cloth is keeping shape from the source image. Also, your target cloth lengths are a bit smaller than the cloth in source person images, so the try-on results have some blur areas in the pants.
@minar09, I can share more results if that help me find the area (pose, segmentation, cloth, cloth mask) of improvement. Also how can I make sure that the target cloth and input image be of same length? here
Maybe you need to crop the person image with a bit longer pants area. Also, I am afraid the CP-VTON+ does not have much improvements to offer at this point. If you want more better quality, then I think you can try with the ACGPN (CVPR 2020). Their virtual try-on quality is state-of-the-art on Viton dataset so far.
@MotiBaadror, did you re train the network with new data format? In some case, the different between training data format and testing data format cause the error.
@thaithanhtuan, I did not retrain the network. I am using pre-trained network. Is it necessary to retrain while using the model for custom data. Although I am using data that is from ecommerce dataset. One more thing that I am using this article to calculate pose_keypoints.
Maybe you need to crop the person image with a bit longer pants area. Also, I am afraid the CP-VTON+ does not have much improvements to offer at this point. If you want more better quality, then I think you can try with the ACGPN (CVPR 2020). Their virtual try-on quality is state-of-the-art on Viton dataset so far.
Hi~I'm appreciate for your efforts in virtual try-on, and what do you think about PF-AFN (CVPR 2021)(https://github.com/geyuying/PF-AFN), would it be better than ACGPN?
Hi @amandazw , actually I haven't tried any new VTON works after 2020. I hope someone else can share their experience with the newer works. Thank you.
I am using custom dataset. More details Batch_size =1 Segmentation Mask from CHIP_PGN results. Pose from openpose Coco Model and changed the code that described here Cloth and cloth mask are from the preprocessed data