def read_data_train(dataset_path, debug=False): h, w = 2048, 2048 dataset = { 'vid_name': [], 'frame_id': [], 'joints3D': [], 'joints2D': [], 'bbox': [], 'img_name': [], 'features': [], } model = spin.get_pretrained_hmr() # training data user_list = range(1, 9) seq_list = range(1, 3) vid_list = list(range(3)) + list(range(4, 9)) # product = product(user_list, seq_list, vid_list) # user_i, seq_i, vid_i = product[process_id] for user_i in user_list: for seq_i in seq_list: seq_path = os.path.join(dataset_path, 'S' + str(user_i), 'Seq' + str(seq_i)) # mat file with annotations annot_file = os.path.join(seq_path, 'annot.mat') annot2 = sio.loadmat(annot_file)['annot2'] annot3 = sio.loadmat(annot_file)['annot3'] # calibration file and camera parameters for j, vid_i in enumerate(vid_list): # image folder imgs_path = os.path.join(seq_path, 'video_' + str(vid_i)) # per frame pattern = os.path.join(imgs_path, '*.jpg') img_list = sorted(glob.glob(pattern)) vid_used_frames = [] vid_used_joints = [] vid_used_bbox = [] vid_segments = [] vid_uniq_id = "subj" + str(user_i) + '_seq' + str( seq_i) + "_vid" + str(vid_i) + "_seg0" for i, img_i in tqdm_enumerate(img_list): # for each image we store the relevant annotations img_name = img_i.split('/')[-1] joints_2d_raw = np.reshape(annot2[vid_i][0][i], (1, 28, 2)) joints_2d_raw = np.append(joints_2d_raw, np.ones((1, 28, 1)), axis=2) joints_2d = convert_kps(joints_2d_raw, "mpii3d", "spin").reshape((-1, 3)) # visualize = True # if visualize == True and i == 500: # import matplotlib.pyplot as plt # # frame = cv2.cvtColor(cv2.imread(img_i), cv2.COLOR_BGR2RGB) # # for k in range(49): # kp = joints_2d[k] # # frame = cv2.circle( # frame.copy(), # (int(kp[0]), int(kp[1])), # thickness=3, # color=(255, 0, 0), # radius=5, # ) # # cv2.putText(frame, f'{k}', (int(kp[0]), int(kp[1]) + 1), cv2.FONT_HERSHEY_SIMPLEX, 1.5, # (0, 255, 0), # thickness=3) # # plt.imshow(frame) # plt.show() joints_3d_raw = np.reshape(annot3[vid_i][0][i], (1, 28, 3)) / 1000 joints_3d = convert_kps(joints_3d_raw, "mpii3d", "spin").reshape((-1, 3)) bbox = get_bbox_from_kp2d( joints_2d[~np.all(joints_2d == 0, axis=1)]).reshape(4) joints_3d = joints_3d - joints_3d[39] # 4 is the root # check that all joints are visible x_in = np.logical_and(joints_2d[:, 0] < w, joints_2d[:, 0] >= 0) y_in = np.logical_and(joints_2d[:, 1] < h, joints_2d[:, 1] >= 0) ok_pts = np.logical_and(x_in, y_in) if np.sum(ok_pts) < joints_2d.shape[0]: vid_uniq_id = "_".join(vid_uniq_id.split("_")[:-1])+ "_seg" +\ str(int(dataset['vid_name'][-1].split("_")[-1][3:])+1) continue dataset['vid_name'].append(vid_uniq_id) dataset['frame_id'].append(img_name.split(".")[0]) dataset['img_name'].append(img_i) dataset['joints2D'].append(joints_2d) dataset['joints3D'].append(joints_3d) dataset['bbox'].append(bbox) vid_segments.append(vid_uniq_id) vid_used_frames.append(img_i) vid_used_joints.append(joints_2d) vid_used_bbox.append(bbox) vid_segments = np.array(vid_segments) ids = np.zeros((len(set(vid_segments)) + 1)) ids[-1] = len(vid_used_frames) + 1 if (np.where( vid_segments[:-1] != vid_segments[1:])[0]).size != 0: ids[1:-1] = (np.where( vid_segments[:-1] != vid_segments[1:])[0]) + 1 # for i in tqdm(range(len(set(vid_segments)))): # features = extract_features(model, np.array(vid_used_frames)[int(ids[i]):int(ids[i+1])], # vid_used_bbox[int(ids[i]):int((ids[i+1]))], # kp_2d=np.array(vid_used_joints)[int(ids[i]):int(ids[i+1])], # dataset='spin', debug=False) # dataset['features'].append(features) for k in dataset.keys(): dataset[k] = np.array(dataset[k]) # dataset['features'] = np.concatenate(dataset['features']) return dataset
def read_test_data(dataset_path): dataset = { 'vid_name': [], 'frame_id': [], 'joints3D': [], 'joints2D': [], 'bbox': [], 'img_name': [], 'features': [], "valid_i": [] } model = spin.get_pretrained_hmr() user_list = range(1, 7) for user_i in user_list: print('Subject', user_i) seq_path = os.path.join(dataset_path, 'mpi_inf_3dhp_test_set', 'TS' + str(user_i)) # mat file with annotations annot_file = os.path.join(seq_path, 'annot_data.mat') mat_as_h5 = h5py.File(annot_file, 'r') annot2 = np.array(mat_as_h5['annot2']) annot3 = np.array(mat_as_h5['univ_annot3']) valid = np.array(mat_as_h5['valid_frame']) vid_used_frames = [] vid_used_joints = [] vid_used_bbox = [] vid_segments = [] vid_uniq_id = "subj" + str(user_i) + "_seg0" for frame_i, valid_i in tqdm(enumerate(valid)): img_i = os.path.join('mpi_inf_3dhp_test_set', 'TS' + str(user_i), 'imageSequence', 'img_' + str(frame_i + 1).zfill(6) + '.jpg') joints_2d_raw = np.expand_dims(annot2[frame_i, 0, :, :], axis=0) joints_2d_raw = np.append(joints_2d_raw, np.ones((1, 17, 1)), axis=2) joints_2d = convert_kps(joints_2d_raw, src="mpii3d_test", dst="spin").reshape((-1, 3)) # visualize = True # if visualize == True: # import matplotlib.pyplot as plt # # frame = cv2.cvtColor(cv2.imread(os.path.join(dataset_path, img_i)), cv2.COLOR_BGR2RGB) # # for k in range(49): # kp = joints_2d[k] # # frame = cv2.circle( # frame.copy(), # (int(kp[0]), int(kp[1])), # thickness=3, # color=(255, 0, 0), # radius=5, # ) # # cv2.putText(frame, f'{k}', (int(kp[0]), int(kp[1]) + 1), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 0), # thickness=3) # # plt.imshow(frame) # plt.show() joints_3d_raw = np.reshape(annot3[frame_i, 0, :, :], (1, 17, 3)) / 1000 joints_3d = convert_kps(joints_3d_raw, "mpii3d_test", "spin").reshape((-1, 3)) joints_3d = joints_3d - joints_3d[ 39] # substract pelvis zero is the root for test bbox = get_bbox_from_kp2d( joints_2d[~np.all(joints_2d == 0, axis=1)]).reshape(4) # check that all joints are visible img_file = os.path.join(dataset_path, img_i) I = cv2.imread(img_file) h, w, _ = I.shape x_in = np.logical_and(joints_2d[:, 0] < w, joints_2d[:, 0] >= 0) y_in = np.logical_and(joints_2d[:, 1] < h, joints_2d[:, 1] >= 0) ok_pts = np.logical_and(x_in, y_in) if np.sum(ok_pts) < joints_2d.shape[0]: vid_uniq_id = "_".join(vid_uniq_id.split("_")[:-1]) + "_seg" + \ str(int(dataset['vid_name'][-1].split("_")[-1][3:]) + 1) continue dataset['vid_name'].append(vid_uniq_id) dataset['frame_id'].append(img_file.split("/")[-1].split(".")[0]) dataset['img_name'].append(img_file) dataset['joints2D'].append(joints_2d) dataset['joints3D'].append(joints_3d) dataset['bbox'].append(bbox) dataset['valid_i'].append(valid_i) vid_segments.append(vid_uniq_id) vid_used_frames.append(img_file) vid_used_joints.append(joints_2d) vid_used_bbox.append(bbox) vid_segments = np.array(vid_segments) ids = np.zeros((len(set(vid_segments)) + 1)) ids[-1] = len(vid_used_frames) + 1 if (np.where(vid_segments[:-1] != vid_segments[1:])[0]).size != 0: ids[1:-1] = (np.where( vid_segments[:-1] != vid_segments[1:])[0]) + 1 # for i in tqdm(range(len(set(vid_segments)))): # features = extract_features(model, np.array(vid_used_frames)[int(ids[i]):int(ids[i + 1])], # vid_used_bbox[int(ids[i]):int(ids[i + 1])], # kp_2d=np.array(vid_used_joints)[int(ids[i]):int(ids[i + 1])], # dataset='spin', debug=False) # dataset['features'].append(features) for k in dataset.keys(): dataset[k] = np.array(dataset[k]) # dataset['features'] = np.concatenate(dataset['features']) return dataset