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