minar09 / cp-vton-plus

Official implementation for "CP-VTON+: Clothing Shape and Texture Preserving Image-Based Virtual Try-On", CVPRW 2020
https://minar09.github.io/cpvtonplus/
MIT License
352 stars 121 forks source link

Results are not good #23

Closed MotiBaadror closed 4 years ago

MotiBaadror commented 4 years ago

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

2940 2978 3700

minar09 commented 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.

MotiBaadror commented 4 years ago

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()
minar09 commented 4 years ago

@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.

MotiBaadror commented 4 years ago

@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

minar09 commented 4 years ago

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.

thaithanhtuan commented 4 years ago

@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.

MotiBaadror commented 4 years ago

@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.

amandazw commented 3 years ago

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?

minar09 commented 3 years ago

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.