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