def read_flow(self, sid1, sid2):
'''
Output:
flow_2to1: (h,w,2) correspondence from image 2 to image 1. corr_2to1[y,x] = [u,v], means image2[y,x] -> image1[v,u]
vis_2: (h,w) visibility mask of image 2.
0: human pixel with correspondence
1: human pixel without correspondece
2: background pixel
'''
fn = os.path.join(self.corr_dir, '%s_%s.corr' % (sid2, sid1))
corr_2to1, vis_2 = flow_util.read_corr(fn)
vis_2 = vis_2[..., np.newaxis]
flow_2to1 = flow_util.corr_to_flow(corr_2to1, vis_2, order='HWC')
if self.opt.vis_smooth_rate > 0:
vis_2b = cv2.medianBlur(vis_2,
self.opt.vis_smooth_rate)[..., np.newaxis]
m = (vis_2 < 2).astype(np.uint8)
vis_2 = vis_2b * m + vis_2 * (1 - m)
return flow_2to1, vis_2
def __getitem__(self, index):
sid1, sid2 = self.id_list[index]
######################
# load data
######################
img_1 = self.read_image(sid1)
img_2 = self.read_image(sid2)
seg_label_1 = self.read_seg(sid1)
seg_label_2 = self.read_seg(sid2)
joint_c_1 = np.array(self.pose_label[sid1])
joint_c_2 = np.array(self.pose_label[sid2])
corr_2to1, vis_2 = self.read_corr(sid1, sid2)
h, w = self.opt.image_size
######################
# augmentation
######################
use_augmentation = self.opt.use_augmentation and self.opt.is_train and self.split == 'train'
if use_augmentation:
# apply random shift and scale on img_2
h, w = self.opt.image_size
dx = np.random.randint(-self.opt.aug_shiftx_range,
self.opt.aug_shiftx_range
) if self.opt.aug_shiftx_range > 0 else 0
dy = np.random.randint(-self.opt.aug_shifty_range,
self.opt.aug_shifty_range
) if self.opt.aug_shifty_range > 0 else 0
sc = self.opt.aug_scale_range**(np.random.rand() * 2 - 1)
M = np.array([[sc, 0, 0.5 * h * (1 - sc) + dx],
[0, sc, 0.5 * w * (1 - sc) + dy]])
img_2 = cv2.warpAffine(img_2,
M,
dsize=(w, h),
flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_REPLICATE)
seg_label_2 = cv2.warpAffine(
seg_label_2,
M,
dsize(w, h),
flags=cv2.INTER_NEAREST,
borderMode=cv2.BORDER_REPLICATE)[..., np.newaxis]
corr_2to1 = cv2.warpAffine(corr_2to1,
M,
dsize=(w, h),
flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_REPLICATE)
vis_2 = cv2.warpAffine(vis_2,
M,
dsize=(w, h),
flags=cv2.INTER_NEAREST,
borderMode=cv2.BORDER_REPLICATE)
v = (d[:, 0] >= 0) & (d[:, 1] >= 0) & (d[:, 0] < w) & (d[:, 1] < h)
jc = joint_c_2.dot(M[:, 0:2].T) + M[:, 2:].T
v_t = (jc[:, 0] >= 0) & (jc[:, 1] >= 0) & (jc[:, 0] <
w) & (jc[:, 1] < h)
v_t = v_t & v
jc[~v_t, :] = -1
joint_c_2 = jc
######################
# pack output data
######################
joint_1 = kp_to_map(img_sz=(w, h),
kps=joint_c_1,
mode=self.opt.joint_mode,
radius=self.opt.joint_radius)
joint_2 = kp_to_map(img_sz=(w, h),
kps=joint_c_2,
mode=self.opt.joint_mode,
radius=self.opt.joint_radius)
seg_1 = seg_label_to_map(seg_label_1, nc=self.opt.seg_nc)
seg_2 = seg_label_to_map(seg_label_2, nc=self.opt.seg_nc)
flow_2to1 = flow_util.corr_to_flow(corr_2to1, vis_2, order='HWC')
flow_2to1[..., 0] = flow_2to1[..., 0].clip(-w, w)
flow_2to1[..., 1] = flow_2to1[..., 1].clip(-h, h)
data = {
'img_1': self.tensor_normalize_std(self.to_tensor(img_1)),
'img_2': self.tensor_normalize_std(self.to_tensor(img_2)),
'joint_c_1': torch.from_numpy(joint_c_1).float(),
'joint_c_2': torch.from_numpy(joint_c_2).float(),
'joint_1': self.to_tensor(joint_1),
'joint_2': self.to_tensor(joint_2),
'seg_label_1': self.to_tensor(seg_label_1),
'seg_label_2': self.to_tensor(seg_label_2),
'seg_1': self.to_tensor(seg_1),
'seg_2': self.to_tensor(seg_2),
'flow_2to1': self.to_tensor(flow_2to1),
'vis_2': self.to_tensor(vis_2),
'id_1': sid1,
'id_2': sid2
}
######################
# color jit
######################
if use_augmentation and self.opt.aug_color_jit:
data['img_1'], data['img_2'] = self.color_jit(
data['img_1'], data['img_2'])
return data
def __getitem__(self, index):
######################
# load data
######################
sid1, sid2 = self.id_list[index]
img_1 = self.read_image(sid1)
img_2 = self.read_image(sid2)
seg_label_1 = self.read_seg(sid1)
seg_label_2 = self.read_seg(sid2)
joint_c_1 = np.array(self.pose_label[sid1])
joint_c_2 = np.array(self.pose_label[sid2])
corr_2to1, vis_2 = self.read_corr(sid1, sid2)
flow_2to1 = flow_util.corr_to_flow(corr_2to1, vis_2, order='HWC')
######################
# augmentation
######################
if self.opt.is_train and self.split == 'train':
# flow_2to1, vis_2, img_2, seg_2, joint_c_2 = self.random_scale_and_shift(flow_2to1, vis_2, img_2, seg_2, joint_c_2, silh_2)
dx = np.random.randint(-self.opt.aug_shiftx_range, self.opt.aug_shiftx_range)
dy = np.random.randint(-self.opt.aug_shifty_range, self.opt.aug_shifty_range)
flow_2to1[...,0] -= dx
flow_2to1[...,1] -= dy
vis_2 = self.shift_image(vis_2, dx, dy, pad=2)
flow_2to1 = self.shift_image(flow_2to1, dx, dy, pad=0) * (vis_2<2).astype(np.float32)
img_2 = self.shift_image(img_2, dx, dy, pad=0)
seg_label_2 = self.shift_image(seg_label_2, dx, dy, pad=0)
v = (joint_c_2[:,0] >= 0) & (joint_c_2[:,1] >= 0)
joint_c_2[v,0] += dx
joint_c_2[v,1] += dy
v = (joint_c_2[:,0]<0)|(joint_c_2[:,0]>self.opt.image_size[1])|(joint_c_2[:,1]<0)|(joint_c_2[:,1]>self.opt.image_size[0])
joint_c_2[v,:]=-1
######################
# create pose representation
######################
joint_1 = kp_to_map(img_sz=(img_1.shape[1], img_1.shape[0]), kps=joint_c_1, mode=self.opt.joint_mode, radius=self.opt.joint_radius)
joint_2 = kp_to_map(img_sz=(img_2.shape[1], img_2.shape[0]), kps=joint_c_2, mode=self.opt.joint_mode, radius=self.opt.joint_radius)
######################
# create seg representation
######################
seg_1 = seg_label_to_map(seg_label_1, nc=self.opt.seg_nc)
seg_2 = seg_label_to_map(seg_label_2, nc=self.opt.seg_nc)
######################
# pack output
######################
data = {
'img_1': self.tensor_normalize_std(self.to_tensor(img_1)),
'img_2': self.tensor_normalize_std(self.to_tensor(img_2)),
'joint_c_1': torch.from_numpy(joint_c_1).float(),
'joint_c_2': torch.from_numpy(joint_c_2).float(),
'joint_1': self.to_tensor(joint_1),
'joint_2': self.to_tensor(joint_2),
'seg_label_1': self.to_tensor(seg_label_1),
'seg_label_2': self.to_tensor(seg_label_2),
'seg_1': self.to_tensor(seg_1),
'seg_2': self.to_tensor(seg_2),
'flow_2to1': self.to_tensor(flow_2to1),
'vis_2': self.to_tensor(vis_2),
'id_1': sid1,
'id_2': sid2
}
return data