Beispiel #1
0
    def __getitem__(self, index):
        t0 = time.time()
        try:
            img_source = self.img_list[index].split('\t')[0]
            target_pose = self.img_list[index].split('\t')[1]
            cloth_img = self.img_list[index].split('\t')[2]
            data_suffix = self.img_list[index].split('\t')[3]
        except:
            img_source = self.img_list[index].split(' ')[0]
            target_pose = self.img_list[index].split(' ')[1]
            cloth_img = self.img_list[index].split(' ')[2]
            data_suffix = self.img_list[index].split(' ')[3]

        data_suffix = data_suffix if data_suffix == 'train' else 'val'
        source_splitext = os.path.splitext(img_source)[0]
        cloth_splitext = os.path.splitext(cloth_img)[0]

        if self.opt.warp_cloth:
            source_splitext = os.path.join(data_suffix, source_splitext)
            cloth_img = os.path.join(data_suffix, cloth_img)
            cloth_splitext = os.path.join(data_suffix, cloth_splitext)

        source_img_path = os.path.join('dataset/images',
                                       source_splitext + '.jpg')
        cloth_parse_path = os.path.join('dataset/cloth_mask',
                                        cloth_splitext + '_mask.png')

        if self.opt.warp_cloth:
            cloth_img_path = os.path.join('dataset/images', cloth_img)
        else:
            cloth_img_path = os.path.join('dataset/cloth_image', cloth_img)

        ### image
        source_img = self.open_transform(source_img_path, False)
        cloth_img = self.open_transform(cloth_img_path, False)
        cloth_parse = self.parse_cloth(cloth_parse_path)
        # parsing
        if self.opt.warp_cloth:
            source_splitext = os.path.join(
                source_splitext.split('/')[0],
                source_splitext.split('/')[2])

        source_parse_vis_path = os.path.join('dataset/parse_cihp',
                                             source_splitext + '_vis.png')
        source_parse_vis = self.transforms['3'](
            Image.open(source_parse_vis_path))
        source_parse_path = os.path.join('dataset/parse_cihp',
                                         source_splitext + '.png')
        source_parse = pose_utils.parsing_embedding(source_parse_path)

        source_parse_shape = np.array(Image.open(source_parse_path))
        source_parse_shape = (source_parse_shape > 0).astype(np.float32)
        source_parse_shape = Image.fromarray(
            (source_parse_shape * 255).astype(np.uint8))
        source_parse_shape = source_parse_shape.resize(
            (self.size[1] // 16, self.size[0] // 16),
            Image.BILINEAR)  # downsample and then upsample
        source_parse_shape = source_parse_shape.resize(
            (self.size[1], self.size[0]), Image.BILINEAR)
        source_parse_shape = self.transforms['1'](source_parse_shape)  # [-1,1]

        source_parse_head = (np.array(Image.open(source_parse_path)) == 1).astype(np.float32) + \
                    (np.array(Image.open(source_parse_path)) == 2).astype(np.float32) + \
                    (np.array(Image.open(source_parse_path)) == 4).astype(np.float32) + \
                    (np.array(Image.open(source_parse_path)) == 13).astype(np.float32)

        # prepare for warped cloth
        phead = torch.from_numpy(source_parse_head)  # [0,1]
        im_h = source_img * phead - (
            1 - phead
        )  # [-1,1], fill -1 for other parts, thus become black visual

        # pose heatmap embedding
        source_pose_path = os.path.join('dataset/pose_coco',
                                        source_splitext + '_keypoints.json')
        with open(source_pose_path, 'r') as f:
            a = json.load(f)
            source_pose = a['people'][0]['pose_keypoints_2d']
        source_pose_loc = pose_utils.pose2loc(source_pose)
        source_pose_embedding = pose_utils.heatmap_embedding(
            self.size, source_pose_loc)

        if self.opt.warp_cloth:
            target_splitext = os.path.splitext(target_pose)[0]
            target_pose_path = os.path.join(
                'dataset/pose_coco', data_suffix,
                target_splitext.split('/')[1] + '_keypoints.json')
            warped_cloth_name = source_splitext.split('/')[0] + '/' + \
                    source_splitext.split('/')[1] + '_' + \
                    target_splitext.split('/')[1] + '_' + \
                    cloth_splitext.split('/')[2] + '_warped_cloth.jpg'
            warped_cloth_name = warped_cloth_name.split(
                '/')[0] + '/' + warped_cloth_name.split('/')[1].split(
                    '-')[0] + '/' + warped_cloth_name.split('/')[1]
        else:
            target_pose_path = os.path.join('dataset/pose_coco', target_pose)
            warped_cloth_name = cloth_splitext

        with open(target_pose_path, 'r') as f:
            a = json.load(f)
            target_pose = a['people'][0]['pose_keypoints_2d']
        target_pose_loc = pose_utils.pose2loc(target_pose)
        target_pose_embedding = pose_utils.heatmap_embedding(
            self.size, target_pose_loc)
        target_pose_img, _ = pose_utils.draw_pose_from_cords(
            target_pose_loc, (256, 192))
        result = {
            'source_parse': source_parse,
            'source_parse_vis': source_parse_vis,
            'source_pose_embedding': source_pose_embedding,
            'target_pose_embedding': target_pose_embedding,
            'target_pose_loc': target_pose_loc,
            'source_image': source_img,
            'cloth_image': cloth_img,
            'cloth_parse': cloth_parse,
            'source_parse_shape': source_parse_shape,
            'im_h': im_h,  # source image head and hair
            'source_image_name': source_splitext,
            'cloth_image_name': cloth_splitext,
            'source_img_path': source_img_path,
            'target_pose_path': target_pose_path,
            'source_parse_vis_path': source_parse_vis_path,
            'target_pose_img': target_pose_img,
            'warped_cloth_name': warped_cloth_name
        }

        return result
Beispiel #2
0
    def __getitem__(self, index):
        t0 = time.time()
        img_ext = '.jpg'
        try:
            img_source = self.img_list[index].split('\t')[0]
            img_target = self.img_list[index].split('\t')[1]
            cloth_img = self.img_list[index].split('\t')[2]
        except:
            img_source = self.img_list[index].split(' ')[0]
            img_target = self.img_list[index].split(' ')[1]
            cloth_img = self.img_list[index].split(' ')[2]

        source_splitext = os.path.splitext(img_source)[0]
        target_splitext = os.path.splitext(img_target)[0]
        cloth_splitext = os.path.splitext(cloth_img)[0]

        # png or jpg
        source_img_path = os.path.join('datasets/zalando/source_images',
                                       source_splitext + img_ext)
        target_img_path = os.path.join('datasets/zalando/target_images',
                                       target_splitext + img_ext)
        cloth_img_path = os.path.join('datasets/zalando/cloth', cloth_img)
        cloth_parse_path = os.path.join('datasets/zalando/cloth_mask',
                                        cloth_splitext + img_ext)
        warped_cloth_name = cloth_img

        # image
        warped_cloth_path = os.path.join('dataset', 'cloth', warped_cloth_name)
        source_img = self.open_transform(source_img_path, False)
        target_img = self.open_transform(target_img_path, False)
        cloth_img = self.open_transform(cloth_img_path, False)
        cloth_parse = self.parse_cloth(cloth_parse_path)

        try:
            warped_cloth_parse_name = cloth_img
            # mask
            warped_cloth_parse_path = os.path.join('dataset', 'cloth_mask',
                                                   warped_cloth_parse_name)
            warped_cloth_parse = self.parse_cloth(warped_cloth_parse_path)
        except:
            warped_cloth_parse = torch.ones(1, 256, 192)

        if os.path.exists(warped_cloth_path):
            warped_cloth_img = self.open_transform(warped_cloth_path, False)
        else:
            warped_cloth_img = cloth_img

        # parsing
        source_parse_vis_path = os.path.join(
            'datasets/zalando/parse_cihp_source',
            source_splitext + '_vis' + '.png')
        target_parse_vis_path = os.path.join(
            'datasets/zalando/parse_cihp_target',
            target_splitext + '_vis' + '.png')
        source_parse_vis = self.transforms['2'](
            Image.open(source_parse_vis_path))
        target_parse_vis = self.transforms['2'](
            Image.open(target_parse_vis_path))

        source_parse_path = os.path.join('datasets/zalando/parse_cihp_source',
                                         source_splitext + '.png')
        target_parse_path = os.path.join('datasets/zalando/parse_cihp_target',
                                         target_splitext + '.png')

        source_parse = pose_utils.parsing_embedding(source_parse_path)
        source_parse_tformed = self.custom_transform(source_parse, True)
        source_parse = torch.from_numpy(source_parse)

        target_parse = pose_utils.parsing_embedding(target_parse_path)
        target_parse_tformed = self.custom_transform(target_parse, True)
        target_parse = torch.from_numpy(target_parse)

        source_parse_shape = np.array(Image.open(source_parse_path))
        source_parse_shape = (source_parse_shape > 0).astype(np.float32)
        source_parse_shape = Image.fromarray(
            (source_parse_shape * 255).astype(np.uint8))
        source_parse_shape = source_parse_shape.resize(
            (self.size[1] // 16, self.size[0] // 16),
            Image.BILINEAR)  # downsample and then upsample
        source_parse_shape = source_parse_shape.resize(
            (self.size[1], self.size[0]), Image.BILINEAR)
        source_parse_shape = self.transforms['2'](source_parse_shape)  # [-1,1]

        source_parse_head = (np.array(Image.open(source_parse_path)) == 1).astype(np.float32) + \
            (np.array(Image.open(source_parse_path)) == 2).astype(np.float32) + \
            (np.array(Image.open(source_parse_path)) == 4).astype(np.float32) + \
            (np.array(Image.open(source_parse_path)) == 13).astype(np.float32)

        target_parse_cloth = (np.array(Image.open(target_parse_path)) == 5).astype(np.float32) + \
            (np.array(Image.open(target_parse_path)) == 6).astype(np.float32) + \
            (np.array(Image.open(target_parse_path)) == 7).astype(np.float32)
        target_parse_cloth = torch.from_numpy(target_parse_cloth)
        # prepare for warped cloth

        phead = torch.from_numpy(source_parse_head)  # [0,1]
        pcm = target_parse_cloth  # [0,1]
        im = target_img  # [-1,1]
        im_c = im * pcm + (
            1 - pcm)  # [-1,1], fill 1 for other parts --> white same as GT ...
        im_h = source_img * phead - (
            1 - phead
        )  # [-1,1], fill -1 for other parts, thus become black visual

        # pose heatmap embedding

        # source_pose_path = os.path.join('datasets/zalando/source_keypoints',
        #                                 source_splitext + '_keypoints.npy')
        # source_pose_org = np.load(source_pose_path)
        # source_pose = []
        # for i in source_pose_org:
        #     source_pose.extend(i)
        # source_pose_loc = pose_utils.pose2loc(source_pose)
        # source_pose_embedding = pose_utils.heatmap_embedding(
        #     self.size, source_pose_loc)

        # target_pose_path = os.path.join('datasets/zalando/target_keypoints',
        #                                 target_splitext + '_keypoints.npy')
        # target_pose_org = np.load(target_pose_path)
        # target_pose = []
        # for i in target_pose_org:
        #     target_pose.extend(i)
        # target_pose_loc = pose_utils.pose2loc(target_pose)
        # target_pose_embedding = pose_utils.heatmap_embedding(
        #     self.size, target_pose_loc)
        # target_pose_img, _ = pose_utils.draw_pose_from_cords(
        #     target_pose_loc, (256, 192))

        # pose heatmap embedding
        source_pose_path = os.path.join('datasets/zalando/source_keypoints',
                                        source_splitext + '_keypoints.json')
        with open(source_pose_path, 'r') as f:
            a = json.load(f)
            source_pose = a['people'][0]['pose_keypoints_2d']
        source_pose_loc = pose_utils.pose2loc(source_pose)
        source_pose_embedding = torch.from_numpy(
            pose_utils.heatmap_embedding(self.size, source_pose_loc))

        target_pose_path = os.path.join('datasets/zalando/target_keypoints',
                                        target_splitext + '_keypoints.json')
        with open(target_pose_path, 'r') as f:
            a = json.load(f)
            target_pose = a['people'][0]['pose_keypoints_2d']
        target_pose_loc = pose_utils.pose2loc(target_pose)
        target_pose_embedding = torch.from_numpy(
            pose_utils.heatmap_embedding(self.size, target_pose_loc))
        # target_pose_img, _ = pose_utils.draw_pose_from_cords(
        #     target_pose_loc, (256, 192))

        # Densepose preprocess source
        source_densepose_path = os.path.join(
            'datasets/zalando/densepose_numpy_source/',
            source_splitext + '.npy')
        source_densepose_data = np.load(source_densepose_path).astype(
            'uint8')  # (256,192,3)
        source_densepose_parts_embeddings = self.parsing_embedding(
            source_densepose_data[:, :, 0])
        source_densepose_parts_embeddings = np.transpose(
            source_densepose_parts_embeddings, axes=(1, 2, 0))
        source_densepose_data_final = np.concatenate(
            (source_densepose_parts_embeddings, source_densepose_data[:, :,
                                                                      1:]),
            axis=-1)  # channel(27), H, W
        source_densepose_data_final = torch.from_numpy(
            np.transpose(source_densepose_data_final, axes=(2, 0, 1)))

        # Densepose preprocess target
        target_densepose_path = os.path.join(
            'datasets/zalando/densepose_numpy_target/',
            source_splitext + '.npy')
        target_densepose_data = np.load(target_densepose_path).astype(
            'uint8')  # (256,192,3)
        target_densepose_parts_embeddings = self.parsing_embedding(
            target_densepose_data[:, :, 0])
        target_densepose_parts_embeddings = np.transpose(
            target_densepose_parts_embeddings, axes=(1, 2, 0))
        target_densepose_data_final = np.concatenate(
            (target_densepose_parts_embeddings, target_densepose_data[:, :,
                                                                      1:]),
            axis=-1)  # channel(27), H, W
        target_densepose_data_final = torch.from_numpy(
            np.transpose(target_densepose_data_final, axes=(2, 0, 1)))

        result = {
            'source_parse': source_parse,
            'source_parse_tformed': source_parse_tformed,
            'target_parse': target_parse,
            'target_parse_tformed': target_parse_tformed,
            'source_parse_vis': source_parse_vis,
            'target_parse_vis': target_parse_vis,
            'source_pose_embedding': source_pose_embedding,
            'target_pose_embedding': target_pose_embedding,
            'target_pose_loc': target_pose_loc,
            'source_image': source_img,
            'target_image': target_img,
            'cloth_image': cloth_img,
            'cloth_parse': cloth_parse,
            'source_parse_shape': source_parse_shape,
            'im_h': im_h,  # source image head and hair
            'im_c': im_c,  # target_cloth_image_warped
            'source_image_name': source_splitext + img_ext,
            'target_image_name': target_splitext + img_ext,
            'cloth_image_name': cloth_splitext + img_ext,
            'warped_cloth_image': warped_cloth_img,
            'warped_cloth_name': warped_cloth_name,
            'warped_cloth_path': warped_cloth_path,
            'source_img_path': source_img_path,
            'target_img_path': target_img_path,
            'target_pose_path': target_pose_path,
            'target_parse_path': target_parse_path,
            'source_parse_vis_path': source_parse_vis_path,
            'target_parse_vis_path': target_parse_vis_path,
            # 'target_pose_img': target_pose_img,
            'warped_cloth_parse': warped_cloth_parse,
            'target_parse_cloth': target_parse_cloth,
            'source_densepose_path': source_densepose_path,
            'source_densepose_data': source_densepose_data_final,
            'target_densepose_path': target_densepose_path,
            'target_densepose_data': target_densepose_data_final
        }

        return result
    def __getitem__(self, index):
        t0 = time.time()
        try:
            img_source = self.img_list[index].split('\t')[0]
            img_target = self.img_list[index].split('\t')[1]
            cloth_img = self.img_list[index].split('\t')[2]
        except:
            img_source = self.img_list[index].split(' ')[0]
            img_target = self.img_list[index].split(' ')[1]
            cloth_img = self.img_list[index].split(' ')[2]

        source_splitext = os.path.splitext(img_source)[0]
        target_splitext = os.path.splitext(img_target)[0]
        cloth_splitext = os.path.splitext(cloth_img)[0]

        # png or jpg
        source_img_path = os.path.join('dataset/images', self.mode,
                                       source_splitext + '.jpg')
        target_img_path = os.path.join('dataset/images', self.mode,
                                       target_splitext + '.jpg')
        cloth_img_path = os.path.join('dataset/images', self.mode, cloth_img)
        cloth_parse_path = os.path.join('dataset/cloth_mask', self.mode,
                                        cloth_splitext + '_mask.png')
        warped_cloth_name = source_splitext.split('/')[0] + '/' + \
                            source_splitext.split('/')[1] + '_' + \
                            target_splitext.split('/')[1] + '_' + \
                            cloth_splitext.split('/')[1] + '_warped_cloth.jpg'
        ### image
        warped_cloth_path = os.path.join('dataset', 'warped_cloth', self.mode,
                                         warped_cloth_name)
        source_img = self.open_transform(source_img_path, False)
        target_img = self.open_transform(target_img_path, False)
        cloth_img = self.open_transform(cloth_img_path, False)
        cloth_parse = self.parse_cloth(cloth_parse_path)

        try:
            warped_cloth_parse_name = source_splitext.split('/')[0] + '/' + \
                source_splitext.split('/')[1] + '_' + \
                target_splitext.split('/')[1] + '_' + \
                cloth_splitext.split('/')[1] + '_warped_cloth_mask.png'
            ### mask
            warped_cloth_parse_path = os.path.join('dataset',
                                                   'warped_cloth_mask',
                                                   self.mode,
                                                   warped_cloth_parse_name)
            warped_cloth_parse = self.parse_cloth(warped_cloth_parse_path)
        except:
            warped_cloth_parse = torch.ones(1, 256, 192)

        if os.path.exists(warped_cloth_path):
            warped_cloth_img = self.open_transform(warped_cloth_path, False)
        else:
            warped_cloth_img = cloth_img
        # parsing
        source_parse_vis_path = os.path.join(
            'dataset/parse_cihp', self.mode,
            source_splitext.split('/')[1] + '_vis.png')
        target_parse_vis_path = os.path.join(
            'dataset/parse_cihp', self.mode,
            target_splitext.split('/')[1] + '_vis.png')
        source_parse_vis = self.transforms['3'](
            Image.open(source_parse_vis_path))
        target_parse_vis = self.transforms['3'](
            Image.open(target_parse_vis_path))

        source_parse_path = os.path.join(
            'dataset/parse_cihp', self.mode,
            source_splitext.split('/')[1] + '.png')
        target_parse_path = os.path.join(
            'dataset/parse_cihp', self.mode,
            target_splitext.split('/')[1] + '.png')

        source_parse = pose_utils.parsing_embedding(source_parse_path)
        target_parse = pose_utils.parsing_embedding(target_parse_path)

        source_parse_shape = np.array(Image.open(source_parse_path))
        source_parse_shape = (source_parse_shape > 0).astype(np.float32)
        source_parse_shape = Image.fromarray(
            (source_parse_shape * 255).astype(np.uint8))
        source_parse_shape = source_parse_shape.resize(
            (self.size[1] // 16, self.size[0] // 16),
            Image.BILINEAR)  # downsample and then upsample
        source_parse_shape = source_parse_shape.resize(
            (self.size[1], self.size[0]), Image.BILINEAR)
        source_parse_shape = self.transforms['1'](source_parse_shape)  # [-1,1]

        source_parse_head = (np.array(Image.open(source_parse_path)) == 1).astype(np.float32) + \
                    (np.array(Image.open(source_parse_path)) == 2).astype(np.float32) + \
                    (np.array(Image.open(source_parse_path)) == 4).astype(np.float32) + \
                    (np.array(Image.open(source_parse_path)) == 13).astype(np.float32)

        target_parse_cloth = (np.array(Image.open(target_parse_path)) == 5).astype(np.float32) + \
                (np.array(Image.open(target_parse_path)) == 6).astype(np.float32) + \
                (np.array(Image.open(target_parse_path)) == 7).astype(np.float32)

        # prepare for warped cloth
        phead = torch.from_numpy(source_parse_head)  # [0,1]
        pcm = torch.from_numpy(target_parse_cloth)  # [0,1]
        im = target_img  # [-1,1]
        im_c = im * pcm + (
            1 - pcm)  # [-1,1], fill 1 for other parts --> white same as GT ...
        im_h = source_img * phead - (
            1 - phead
        )  # [-1,1], fill -1 for other parts, thus become black visual

        # pose heatmap embedding
        source_pose_path = os.path.join(
            'dataset/pose_coco', self.mode,
            source_splitext.split('/')[1] + '_keypoints.json')
        with open(source_pose_path, 'r') as f:
            a = json.load(f)
            source_pose = a['people'][0]['pose_keypoints_2d']
        source_pose_loc = pose_utils.pose2loc(source_pose)
        source_pose_embedding = pose_utils.heatmap_embedding(
            self.size, source_pose_loc)

        target_pose_path = os.path.join(
            'dataset/pose_coco', self.mode,
            target_splitext.split('/')[1] + '_keypoints.json')
        with open(target_pose_path, 'r') as f:
            a = json.load(f)
            target_pose = a['people'][0]['pose_keypoints_2d']
        target_pose_loc = pose_utils.pose2loc(target_pose)
        target_pose_embedding = pose_utils.heatmap_embedding(
            self.size, target_pose_loc)
        target_pose_img, _ = pose_utils.draw_pose_from_cords(
            target_pose_loc, (256, 192))

        result = {
            'source_parse': source_parse,
            'target_parse': target_parse,
            'source_parse_vis': source_parse_vis,
            'target_parse_vis': target_parse_vis,
            'source_pose_embedding': source_pose_embedding,
            'target_pose_embedding': target_pose_embedding,
            'target_pose_loc': target_pose_loc,
            'source_image': source_img,
            'target_image': target_img,
            'cloth_image': cloth_img,
            'cloth_parse': cloth_parse,
            'source_parse_shape': source_parse_shape,
            'im_h': im_h,  # source image head and hair
            'im_c': im_c,  # target_cloth_image_warped
            'source_image_name': source_splitext,
            'target_image_name': target_splitext,
            'cloth_image_name': cloth_splitext,
            'warped_cloth_image': warped_cloth_img,
            'warped_cloth_name': warped_cloth_name,
            'warped_cloth_path': warped_cloth_path,
            'source_img_path': source_img_path,
            'target_img_path': target_img_path,
            'target_pose_path': target_pose_path,
            'target_parse_path': target_parse_path,
            'source_parse_vis_path': source_parse_vis_path,
            'target_parse_vis_path': target_parse_vis_path,
            'target_pose_img': target_pose_img,
            'warped_cloth_parse': warped_cloth_parse,
            'target_parse_cloth': target_parse_cloth,
        }

        return result