def set_input(self, input):
# move to GPU and change data types
random.shuffle(self.keys)
if len(self.gpu_ids) > 0:
self.input_fullP1 = input['P1'].cuda(self.gpu_ids[0], async=True)
self.input_fullP2 = input['P2'].cuda(self.gpu_ids[0], async=True)
input_P1mask = input['P1masks'].cuda(self.gpu_ids[0], async=True)
input_P2mask = input['P2masks'].cuda(self.gpu_ids[0], async=True)
input_P1mask, _ = pose_utils.obtain_mask(input_P1mask, self.mask_id,
self.keys)
input_P2mask, input_P2back = pose_utils.obtain_mask(
input_P2mask, self.mask_id, self.keys)
self.input_P1 = self.input_fullP1.repeat(3, 1, 1, 1) * input_P1mask
self.input_P2 = self.input_fullP2.repeat(3, 1, 1, 1) * input_P2mask
self.input_BP1 = pose_utils.cords_to_map(input['BP1'],
input['P1masks'],
self.mask_id, self.keys,
self.GPU, self.opt)
self.input_BP2 = pose_utils.cords_to_map(input['BP2'],
input['P2masks'],
self.mask_id, self.keys,
self.GPU, self.opt)
self.image_paths = []
for i in range(self.opt.batchSize):
self.image_paths.append(
os.path.splitext(input['P1_path'][i])[0] + '_2_' +
input['P2_path'][i])
def set_input_test(self, input):
self.input_fullP1 = input['P2'].cuda()
input_P1mask = input['P2masks'].cuda()
self.input_fullP2 = input['P2'].cuda()
input_P2mask = input['P2masks'].cuda()
input_P1mask, input_P1backmask = pose_utils.obtain_mask(
input_P1mask, self.mask_id, self.keys)
input_P2mask, input_P2backmask = pose_utils.obtain_mask(
input_P2mask, self.mask_id, self.keys)
self.input_P1 = self.input_fullP1.repeat(3, 1, 1, 1) * input_P1mask
self.input_P1_back = self.input_fullP1 * input_P1backmask
self.input_P2 = self.input_fullP2.repeat(3, 1, 1, 1) * input_P2mask
self.input_P2mask = input_P2mask.float()
self.input_P1backmask = input_P1backmask
self.input_P2backmask = input_P2backmask
self.input_BP1 = pose_utils.cords_to_map(input['BP1'],
input['P1masks'],
self.mask_id, self.keys,
self.GPU, self.opt,
input['affine'])
self.input_BP2 = pose_utils.cords_to_map(input['BP2'],
input['P2masks'],
self.mask_id, self.keys,
self.GPU, self.opt)
self.image_paths = []
for i in range(self.opt.batchSize):
self.image_paths.append(
os.path.splitext(input['P1_path'][i])[0] + '_2_' +
input['P2_path'][i])
def obtain_bone(self, name, affine_matrix):
string = self.annotation_file.loc[name]
array = pose_utils.load_pose_cords_from_strings(string['keypoints_y'], string['keypoints_x'])
pose = pose_utils.cords_to_map(array, self.load_size, self.opt.old_size, affine_matrix)
pose = np.transpose(pose,(2, 0, 1))
pose = torch.Tensor(pose)
return pose
def set_input(self, input):
# move to GPU and change data types
random.shuffle(self.keys)
if len(self.gpu_ids) > 0:
self.input_fullP1 = input['P1'].cuda()
self.input_fullP2 = input['P2'].cuda()
input_P1mask = input['P1masks'].cuda()
input_P2mask = input['P2masks'].cuda()
res_mask = []
for key in self.target_mask_id.keys():
tmpmask = []
for k in self.target_mask_id[key]:
tmpmask.append(
torch.where(input_P2mask == k,
torch.ones_like(input_P2mask),
torch.zeros_like(input_P2mask)))
tmpmask = torch.stack(tmpmask)
res_mask.append(torch.sum(tmpmask, axis=0))
self.target_layout = torch.stack(res_mask, 1)
self.target_layout = self.target_layout.float()
input_P1mask, _ = pose_utils.obtain_mask(input_P1mask, self.mask_id,
self.keys)
input_P2mask, _ = pose_utils.obtain_mask(input_P2mask, self.mask_id,
self.keys)
self.input_P1 = self.input_fullP1.repeat(3, 1, 1, 1) * input_P1mask
self.input_P2 = self.input_fullP2.repeat(3, 1, 1, 1) * input_P2mask
self.input_BP1 = pose_utils.cords_to_map(input['BP1'],
input['P1masks'],
self.channel_id, self.keys,
self.GPU, self.opt,
input['affine'])
self.input_BP2 = pose_utils.cords_to_map(input['BP2'],
input['P2masks'],
self.channel_id, self.keys,
self.GPU, self.opt)
self.image_paths = []
for i in range(self.opt.batchSize):
self.image_paths.append(
os.path.splitext(input['P1_path'][i])[0] + '_2_' +
input['P2_path'][i])
def __getitem__(self, index):
A_path = self.A_paths[index % self.A_size]
# if self.opt.serial_batches:
if True:
index_B = index % self.B_size
else:
index_B = random.randint(0, self.B_size - 1)
B_path = self.B_paths[index_B]
print(('(A, B) = (%d, %d)' % (index, index_B)))
A_img = Image.open(A_path).convert('RGB')
B_img = Image.open(B_path).convert('RGB')
downsample_A_img = A_img.resize((128, 128), Image.BICUBIC)
downsample_B_img = B_img.resize((128, 128), Image.BICUBIC)
A_img_name = A_path.split('/')[-1]
A_row = self.annotation_file.loc[A_img_name]
A_kp_array = pose_utils.load_pose_cords_from_strings(
A_row['keypoints_y'], A_row['keypoints_x'])
#downsample pose location
downsample_A_kp_array = self.downsample_pose_array(A_kp_array)
A_pose = pose_utils.cords_to_map(A_kp_array, self._image_size)
A_pose = numpy.transpose(
A_pose, (2, 0, 1)) # A_pose: 256 x 256 x 18 => 18 x 256 x 256
A_pose = torch.Tensor(A_pose)
downsample_A_pose = pose_utils.cords_to_map(downsample_A_kp_array,
(128, 128))
downsample_A_pose = numpy.transpose(downsample_A_pose, (2, 0, 1))
downsample_A_pose = torch.Tensor(downsample_A_pose)
B_img_name = B_path.split('/')[-1]
B_row = self.annotation_file.loc[B_img_name]
B_kp_array = pose_utils.load_pose_cords_from_strings(
B_row['keypoints_y'], B_row['keypoints_x'])
#downsample pose location
downsample_B_kp_array = self.downsample_pose_array(B_kp_array)
B_pose = pose_utils.cords_to_map(B_kp_array, (self._image_size))
B_pose = numpy.transpose(
B_pose, (2, 0, 1)) # B_pose: 256 x 256 x 18 => 18 x 256 x 256
B_pose = torch.Tensor(B_pose)
downsample_B_pose = pose_utils.cords_to_map(downsample_B_kp_array,
(128, 128))
downsample_B_pose = numpy.transpose(
downsample_B_pose,
(2, 0, 1)) # B_pose: 256 x 256 x 18 => 18 x 256 x 256
downsample_B_pose = torch.Tensor(downsample_B_pose)
A = self.transform(A_img)
downsample_A = self.transform(downsample_A_img)
downsample_B = self.transform(downsample_B_img)
AtoB_warps, AtoB_masks = self.compute_cord_warp(A_kp_array, B_kp_array)
BtoA_warps, BtoA_masks = self.compute_cord_warp(B_kp_array, A_kp_array)
downsample_AtoB_warps, downsample_AtoB_masks = self.compute_cord_warp(
downsample_A_kp_array, downsample_B_kp_array, img_size=(128, 128))
downsample_BtoA_warps, downsample_BtoA_masks = self.compute_cord_warp(
downsample_B_kp_array, downsample_A_kp_array, img_size=(128, 128))
# Compute the face warp
face_A_warp = self.compute_cord_face_warp(A_kp_array)
face_B_warp = self.compute_cord_face_warp(B_kp_array)
A_parsing_file = './parsing/' + A_img_name.split('.jpg')[0] + '.npy'
# A_parsing_file = '/data/songsijie/deepfashion_parsing/merge_parsing/' + A_img_name.split('.jpg')[0] + '.npy'
A_parsing_data = numpy.load(A_parsing_file)
A_parsing = numpy.zeros((9, 256, 256), dtype='int8')
for id in range(9):
A_parsing[id] = (A_parsing_data == id + 1).astype('int8')
B_parsing_file = './parsing/' + B_img_name.split('.jpg')[0] + '.npy'
# B_parsing_file = '/data/songsijie/deepfashion_parsing/merge_parsing/' + B_img_name.split('.jpg')[0] + '.npy'
B_parsing_data = numpy.load(B_parsing_file)
B_parsing = numpy.zeros((9, 256, 256), dtype='int8')
for id in range(9):
B_parsing[id] = (B_parsing_data == id + 1).astype('int8')
####Downsample A_parsing & B_parsing
# the parsing size became 9 x 128 x 128
[X, Y] = numpy.meshgrid(list(range(0, 256, 2)), list(range(0, 256, 2)))
downsample_A_parsing = A_parsing[:, Y, X]
downsample_B_parsing = B_parsing[:, Y, X]
input_nc = self.opt.input_nc
output_nc = self.opt.output_nc
if input_nc == 1: # RGB to gray
tmp = A[0, ...] * 0.299 + A[1, ...] * 0.587 + A[2, ...] * 0.114
A = tmp.unsqueeze(0)
if output_nc == 1: # RGB to gray
tmp = B[0, ...] * 0.299 + B[1, ...] * 0.587 + B[2, ...] * 0.114
B = tmp.unsqueeze(0)
#A: Image A
#B_pose: B_pose
#A_pose: A_pose
return {
'A': A,
'A_pose': A_pose,
'B_pose': B_pose,
'AtoB_warps': AtoB_warps.astype('float32'),
'BtoA_warps': BtoA_warps.astype('float32'),
'AtoB_masks': AtoB_masks.astype('float32'),
'BtoA_masks': BtoA_masks.astype('float32'),
'A_parsing': A_parsing.astype('float32'),
'B_parsing': B_parsing.astype('float32'),
'A_128': downsample_A,
'B_128': downsample_B,
'A_pose_128': downsample_A_pose,
'B_pose_128': downsample_B_pose,
'AtoB_warps_128': downsample_AtoB_warps.astype('float32'),
'BtoA_warps_128': downsample_BtoA_warps.astype('float32'),
'AtoB_masks_128': downsample_AtoB_masks.astype('float32'),
'BtoA_masks_128': downsample_BtoA_masks.astype('float32'),
'A_parsing_128': downsample_A_parsing.astype('float32'),
'B_parsing_128': downsample_B_parsing.astype('float32'),
'face_A_warp': face_A_warp.astype('float32'),
'face_B_warp': face_B_warp.astype('float32'),
'A_paths': A_path,
'B_paths': B_path,
}