def make_tdpw(): root = '/globalwork/datasets/3DPW' body_joint_names = ( 'pelv,lhip,rhip,bell,lkne,rkne,spin,lank,rank,thor,ltoe,rtoe,neck,lcla,rcla,head,lsho,' 'rsho,lelb,relb,lwri,rwri,lhan,rhan'.split(',')) selected_joints = [*range(1, 24), 0] joint_names = [body_joint_names[j] for j in selected_joints] edges = ('head-neck-lcla-lsho-lelb-lwri-lhan,' 'neck-rcla-rsho-relb-rwri-rhan,' 'neck-thor-spin-bell-pelv-lhip-lkne-lank-ltoe,' 'pelv-rhip-rkne-rank-rtoe') joint_info = p3ds.JointInfo(joint_names, edges) def get_examples(phase, pool): result = [] seq_filepaths = glob.glob(f'{root}/sequenceFiles/{phase}/*.pkl') for filepath in seq_filepaths: with open(filepath, 'rb') as f: seq = pickle.load(f, encoding='latin1') seq_name = seq['sequence'] intrinsics = seq['cam_intrinsics'] extrinsics_per_frame = seq['cam_poses'] for i_person, (coord_seq, coords2d_seq, trans_seq, camvalid_seq) in enumerate(zip( seq['jointPositions'], seq['poses2d'], seq['trans'], seq['campose_valid'])): for i_frame, (coords, coords2d, trans, extrinsics, campose_valid) in enumerate( zip(coord_seq, coords2d_seq, trans_seq, extrinsics_per_frame, camvalid_seq)): if not campose_valid or np.all(coords2d == 0): continue impath = f'{root}/imageFiles/{seq_name}/image_{i_frame:05d}.jpg' camera = cameralib.Camera( extrinsic_matrix=extrinsics, intrinsic_matrix=intrinsics, world_up=(0, 1, 0)) camera.t *= 1000 world_coords = (coords.reshape(-1, 3))[selected_joints] * 1000 camera2 = cameralib.Camera(intrinsic_matrix=intrinsics, world_up=(0, -1, 0)) camcoords = camera.world_to_camera(world_coords) imcoords = camera.world_to_image(world_coords) bbox = boxlib.expand(boxlib.bb_of_points(imcoords), 1.15) ex = p3ds.Pose3DExample(impath, camcoords, bbox=bbox, camera=camera2) noext, ext = os.path.splitext(os.path.relpath(impath, root)) new_image_relpath = f'tdpw_downscaled/{noext}_{i_person:03d}.jpg' pool.apply_async( make_efficient_example, (ex, new_image_relpath, 1, False, "2021-07-09T12:28:07"), callback=result.append) return result with util.BoundedPool(None, 120) as pool: train_examples = get_examples('train', pool) val_examples = get_examples('validation', pool) test_examples = get_examples('test', pool) test_examples = [*train_examples, *val_examples, *test_examples] test_examples.sort(key=lambda ex: ex.image_path) return p3ds.Pose3DDataset(joint_info, None, None, test_examples)
def make_mpii(): joint_names = 'rank,rkne,rhip,lhip,lkne,lank,pelv,thor,neck,head,rwri,relb,rsho,lsho,lelb,lwri' edges = 'lsho-lelb-lwri,rsho-relb-rwri,lhip-lkne-lank,rhip-rkne-rank,neck-head,pelv-thor' joint_info_full = JointInfo(joint_names, edges) joint_names_used = 'rank,rkne,rhip,lhip,lkne,lank,rwri,relb,rsho,lsho,lelb,lwri' joint_info_used = JointInfo(joint_names_used, edges) dataset = Pose2DDataset(joint_info_used) selected_joints = [joint_info_full.ids[name] for name in joint_info_used.names] mat_path = f'{paths.DATA_ROOT}/mpii/mpii_human_pose_v1_u12_1.mat' s = matlabfile.load(mat_path).RELEASE annolist = np.atleast_1d(s.annolist) pool = util.BoundedPool(None, 120) for anno, is_train, rect_ids in zip(annolist, util.progressbar(s.img_train), s.single_person): if not is_train: continue image_path = f'mpii/images/{anno.image.name}' annorect = np.atleast_1d(anno.annorect) rect_ids = np.atleast_1d(rect_ids) - 1 for rect_id in rect_ids: rect = annorect[rect_id] if 'annopoints' not in rect or len(rect.annopoints) == 0: continue coords = np.full( shape=[joint_info_full.n_joints, 2], fill_value=np.nan, dtype=np.float32) for joint in np.atleast_1d(rect.annopoints.point): coords[joint.id] = [joint.x, joint.y] coords = coords[selected_joints] rough_person_center = np.float32([rect.objpos.x, rect.objpos.y]) rough_person_size = rect.scale * 200 # Shift person center down like [Sun et al. 2018], who say this is common on MPII rough_person_center[1] += 0.075 * rough_person_size topleft = np.array(rough_person_center) - np.array(rough_person_size) / 2 bbox = np.array([topleft[0], topleft[1], rough_person_size, rough_person_size]) ex = Pose2DExample(image_path, coords, bbox=bbox) new_im_path = image_path.replace('mpii', 'mpii_downscaled') without_ext, ext = os.path.splitext(new_im_path) new_im_path = f'{without_ext}_{rect_id:02d}{ext}' pool.apply_async( make_efficient_example, (ex, new_im_path), callback=dataset.examples[TRAIN].append) print('Waiting for tasks...') pool.close() pool.join() print('Done...') dataset.examples[TRAIN].sort(key=lambda x: x.image_path) return dataset
def generate_all_visible_boxes(composites, i_valids): """Saves the visible bounding box for each person instance. Visible bounding box means the box around the unoccluded foreground mask. A part of the foreground mask may be occluded by a person or chair pasted in front. """ result = np.zeros((composites.shape[0], 4, 4), np.float32) def setter(i): def fn(val): result[i] = val return fn with util.BoundedPool(None, 120) as pool: for i, sample in zip(i_valids, util.progressbar(composites[i_valids])): pool.apply_async(get_visible_boxes, (sample, ), callback=setter(i)) np.save(f'{paths.DATA_ROOT}/muco/visible_boxes.npy', result)
def make_muco(): joint_info, selected_joints = make_joint_info() root_3dhp = f'{paths.DATA_ROOT}/3dhp' root_muco = f'{paths.DATA_ROOT}/muco' sample_info = np.load(f'{root_muco}/composite_frame_origins.npy') n_all_joints = 28 valid_indices = list(np.load(f'{root_muco}/valid_composite_frame_indices.npy')) all_detections = util.load_pickle(f'{root_muco}/yolov3_detections.pkl') all_detections = np.array([all_detections[k] for k in sorted(all_detections.keys())]) all_visible_boxes = np.load(f'{root_muco}/visible_boxes.npy') matloader = functools.lru_cache(1024)(matlabfile.load) @functools.lru_cache(1024) def get_world_coords(i_subject, i_seq, i_cam, anno_name): seqpath = f'{root_3dhp}/S{i_subject}/Seq{i_seq}' anno_file = matloader(f'{seqpath}/annot.mat') camcoords = anno_file[anno_name][i_cam].reshape( [-1, n_all_joints, 3])[:, selected_joints] camera = load_cameras(f'{seqpath}/camera.calibration')[i_cam] world_coords = [camera.camera_to_world(c) for c in camcoords] return world_coords examples = [] with util.BoundedPool(None, 120) as pool: for i_sample, people, detections, visible_boxes in zip( util.progressbar(valid_indices), sample_info[valid_indices], all_detections[valid_indices], all_visible_boxes[valid_indices]): detections = [box for box in detections if box[-1] > 0.1] if not detections: continue filename = f'{i_sample + 1:06d}.jpg' image_relpath = f'unaugmented_set_001/{filename[:2]}/{filename[:4]}/{filename}' gt_people = [] for i_person, ((i_subject, i_seq, i_cam, i_frame), visible_box) in enumerate( zip(people, visible_boxes)): seqpath = f'{root_3dhp}/S{i_subject}/Seq{i_seq}' world_coords = get_world_coords(i_subject, i_seq, i_cam, 'annot3')[i_frame] univ_world_coords = get_world_coords( i_subject, i_seq, i_cam, 'univ_annot3')[i_frame] camera = load_cameras(f'{seqpath}/camera.calibration')[i_cam] im_coords = camera.world_to_image(world_coords) coord_bbox = boxlib.expand(boxlib.intersect( boxlib.bb_of_points(im_coords), boxlib.full_box([2048, 2048])), 1.05) bbox = boxlib.intersect_vertical(visible_box, coord_bbox) ex = p3ds.Pose3DExample( image_relpath, world_coords, bbox, camera, mask=None, univ_coords=univ_world_coords) gt_people.append(ex) if not gt_people: continue iou_matrix = np.array([[boxlib.iou(gt_person.bbox, box[:4]) for box in detections] for gt_person in gt_people]) gt_indices, det_indices = scipy.optimize.linear_sum_assignment(-iou_matrix) for i_gt, i_det in zip(gt_indices, det_indices): gt_box = gt_people[i_gt].bbox det_box = detections[i_det] if (iou_matrix[i_gt, i_det] > 0.1 and boxlib.area(det_box) < 2 * boxlib.area(gt_box)): ex = gt_people[i_gt] ex.bbox = np.array(detections[i_det][:4]) pool.apply_async(make_efficient_example, (ex, root_muco, i_gt), callback=examples.append) examples.sort(key=lambda ex: ex.image_path) return p3ds.Pose3DDataset(joint_info, examples)
def make_mpi_inf_3dhp(camera_ids=(0, 1, 2, 4, 5, 6, 7, 8)): all_short_names = ( 'spi3,spi4,spi2,spin,pelv,neck,head,htop,lcla,lsho,lelb,lwri,lhan,rcla,rsho,relb,rwri,' 'rhan,lhip,lkne,lank,lfoo,ltoe,rhip,rkne,rank,rfoo,rtoe'.split(',')) test_set_selected_joints = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 15, 16, 14] selected_joints = [7, 5, 14, 15, 16, 9, 10, 11, 23, 24, 25, 18, 19, 20, 3, 6, 4] joint_names = [all_short_names[j] for j in selected_joints] edges = ('htop-head-neck-lsho-lelb-lwri,neck-rsho-relb-rwri,neck-spin-pelv-lhip-lkne-lank,' 'pelv-rhip-rkne-rank') joint_info = p3ds.JointInfo(joint_names, edges) root_3dhp = f'{paths.DATA_ROOT}/3dhp' detections_all = util.load_pickle(f'{paths.DATA_ROOT}/3dhp/yolov3_person_detections.pkl') ################################# # TRAINING AND VALIDATION SET ################################# num_frames = np.asarray( [[6416, 12430], [6502, 6081], [12488, 12283], [6171, 6675], [12820, 12312], [6188, 6145], [6239, 6320], [6468, 6054]]) train_subjects = [0, 1, 2, 3, 4, 5, 6] valid_subjects = [7] # this is my own arbitrary split for validation (Istvan Sarandi) train_examples = [] valid_examples = [] pool = util.BoundedPool(None, 120) for i_subject, i_seq, i_cam in itertools.product( train_subjects + valid_subjects, range(2), camera_ids): seqpath = f'{root_3dhp}/S{i_subject + 1}/Seq{i_seq + 1}' print(f'Processing {seqpath} camera {i_cam}') cam3d_coords = [ann.reshape([ann.shape[0], -1, 3])[:, selected_joints] for ann in matlabfile.load(f'{seqpath}/annot.mat')['annot3']] univ_cam3d_coords = [ann.reshape([ann.shape[0], -1, 3])[:, selected_joints] for ann in matlabfile.load(f'{seqpath}/annot.mat')['univ_annot3']] cameras = load_cameras(f'{seqpath}/camera.calibration') examples_container = train_examples if i_subject in train_subjects else valid_examples frame_step = 5 prev_coords = None camera = cameras[i_cam] n_frames = num_frames[i_subject, i_seq] if i_subject == 5 and i_seq == 1 and i_cam == 2: # This video is shorter for some reason n_frames = 3911 for i_frame in util.progressbar(range(0, n_frames, frame_step)): image_relpath = ( f'3dhp/S{i_subject + 1}/Seq{i_seq + 1}/' f'imageSequence/img_{i_cam}_{i_frame:06d}.jpg') cam_coords = cam3d_coords[i_cam][i_frame] world_coords = cameras[i_cam].camera_to_world(cam_coords) univ_camcoords = univ_cam3d_coords[i_cam][i_frame] univ_world_coords = cameras[i_cam].camera_to_world(univ_camcoords) # Check if the joints are within the image frame bounds if not np.all(camera.is_visible(world_coords, [2048, 2048])): continue im_coords = camera.camera_to_image(cam_coords) bbox = get_bbox(im_coords, image_relpath, detections_all) # Adaptive temporal sampling if (prev_coords is not None and np.all(np.linalg.norm(world_coords - prev_coords, axis=1) < 100)): continue prev_coords = world_coords mask_path = image_relpath.replace('imageSequence', 'FGmasks') new_image_relpath = image_relpath.replace('3dhp', '3dhp_downscaled') ex = p3ds.Pose3DExample( image_relpath, world_coords, bbox, camera, mask=mask_path, univ_coords=univ_world_coords) pool.apply_async(make_efficient_example, (ex, new_image_relpath, 1, True), callback=examples_container.append) print('Waiting for tasks...') pool.close() pool.join() print('Done...') ################################# # TEST SET ################################# test_examples = [] cam1_4 = make_3dhp_test_camera( sensor_size=np.array([10, 10]), im_size=np.array([2048, 2048]), focal_length=7.32506, pixel_aspect=1.00044, center_offset=np.array([-0.0322884, 0.0929296]), distortion=None, origin=np.array([3427.28, 1387.86, 309.42]), up=np.array([-0.208215, 0.976233, 0.06014]), right=np.array([0.000575281, 0.0616098, -0.9981])) cam5_6 = make_3dhp_test_camera( sensor_size=np.array([10, 5.625]), im_size=np.array([1920, 1080]), focal_length=8.770747185, pixel_aspect=0.993236423, center_offset=np.array([-0.104908645, 0.104899704]), distortion=np.array([-0.276859611, 0.131125256, -0.000360494, -0.001149441, -0.049318332]), origin=np.array([-2104.3074, 1038.6707, -4596.6367]), up=np.array([0.025272345, 0.995038509, 0.096227370]), right=np.array([-0.939647257, -0.009210289, 0.342020929])) activity_names = [ 'Stand/Walk', 'Exercise', 'Sit on Chair', 'Reach/Crouch', 'On Floor', 'Sports', 'Misc.'] for i_subject in range(1, 7): seqpath = f'{root_3dhp}/TS{i_subject}' annotation_path = f'{seqpath}/annot_data.mat' with h5py.File(annotation_path, 'r') as m: cam3d_coords = np.array(m['annot3'])[:, 0, test_set_selected_joints] univ_cam3d_coords = np.array(m['univ_annot3'])[:, 0, test_set_selected_joints] valid_frames = np.where(m['valid_frame'][:, 0])[0] activity_ids = m['activity_annotation'][:, 0].astype(int) - 1 camera = cam1_4 if i_subject <= 4 else cam5_6 scene = ['green-screen', 'no-green-screen', 'outdoor'][(i_subject - 1) // 2] for i_frame in valid_frames: image_relpath = f'3dhp/TS{i_subject}/imageSequence/img_{i_frame + 1:06d}.jpg' cam_coords = cam3d_coords[i_frame] univ_camcoords = univ_cam3d_coords[i_frame] activity = activity_names[activity_ids[i_frame]] world_coords = camera.camera_to_world(cam_coords) univ_world_coords = camera.camera_to_world(univ_camcoords) im_coords = camera.camera_to_image(cam_coords) bbox = get_bbox(im_coords, image_relpath, detections_all) ex = p3ds.Pose3DExample( image_relpath, world_coords, bbox, camera, activity_name=activity, scene_name=scene, univ_coords=univ_world_coords) test_examples.append(ex) train_examples.sort(key=lambda x: x.image_path) valid_examples.sort(key=lambda x: x.image_path) test_examples.sort(key=lambda x: x.image_path) return p3ds.Pose3DDataset(joint_info, train_examples, valid_examples, test_examples)
def make_h36m(train_subjects=(1, 5, 6, 7, 8), valid_subjects=(), test_subjects=(9, 11), correct_S9=True, partial_visibility=False): joint_names = ('rhip,rkne,rank,lhip,lkne,lank,tors,neck,head,htop,' 'lsho,lelb,lwri,rsho,relb,rwri,pelv'.split(',')) edges = ('htop-head-neck-lsho-lelb-lwri,neck-rsho-relb-rwri,' 'neck-tors-pelv-lhip-lkne-lank,pelv-rhip-rkne-rank') joint_info = ps3d.JointInfo(joint_names, edges) if not util.all_disjoint(train_subjects, valid_subjects, test_subjects): raise Exception('Set of train, val and test subject must be disjoint.') # use last subject of the non-test subjects for validation train_examples = [] test_examples = [] valid_examples = [] pool = util.BoundedPool(None, 120) if partial_visibility: dir_suffix = '_partial' further_expansion_factor = 1.8 else: dir_suffix = '' if correct_S9 else 'incorrect_S9' further_expansion_factor = 1 for i_subject in [*test_subjects, *train_subjects, *valid_subjects]: if i_subject in train_subjects: examples_container = train_examples elif i_subject in valid_subjects: examples_container = valid_examples else: examples_container = test_examples frame_step = 5 if i_subject in train_subjects else 64 for activity_name, camera_id in itertools.product( get_activity_names(i_subject), range(4)): print(f'Processing S{i_subject} {activity_name} {camera_id}') image_relpaths, world_coords_all, bboxes, camera = get_examples( i_subject, activity_name, camera_id, frame_step=frame_step, correct_S9=correct_S9) prev_coords = None for image_relpath, world_coords, bbox in zip( util.progressbar(image_relpaths), world_coords_all, bboxes): # Using very similar examples is wasteful when training. Therefore: # skip frame if all keypoints are within a distance compared to last stored frame. # This is not done when testing, as it would change the results. if (i_subject in train_subjects and prev_coords is not None and np.all( np.linalg.norm(world_coords - prev_coords, axis=1) < 100)): continue prev_coords = world_coords activity_name = activity_name.split(' ')[0] ex = ps3d.Pose3DExample(image_relpath, world_coords, bbox, camera, activity_name=activity_name) new_image_relpath = image_relpath.replace( 'h36m', f'h36m_downscaled{dir_suffix}') pool.apply_async( make_efficient_example, (ex, new_image_relpath, further_expansion_factor), callback=examples_container.append) print('Waiting for tasks...') pool.close() pool.join() print('Done...') train_examples.sort(key=lambda x: x.image_path) valid_examples.sort(key=lambda x: x.image_path) test_examples.sort(key=lambda x: x.image_path) return ps3d.Pose3DDataset(joint_info, train_examples, valid_examples, test_examples)
def make_h36m(train_subjects=(1, 5, 6, 7, 8), valid_subjects=(), test_subjects=(9, 11), correct_S9=True, partial_visibility=False): joint_names = ('rhip,rkne,rank,lhip,lkne,lank,tors,neck,head,htop,' 'lsho,lelb,lwri,rsho,relb,rwri,pelv'.split(',')) j = ps3d.JointInfo.make_id_map(joint_names) edges = [(j.htop, j.head), (j.head, j.neck), (j.lsho, j.neck), (j.lelb, j.lsho), (j.lwri, j.lelb), (j.rsho, j.neck), (j.relb, j.rsho), (j.rwri, j.relb), (j.neck, j.tors), (j.tors, j.pelv), (j.lhip, j.pelv), (j.lkne, j.lhip), (j.lank, j.lkne), (j.rhip, j.pelv), (j.rkne, j.rhip), (j.rank, j.rkne)] joint_info = ps3d.JointInfo(j, edges) if not util.all_disjoint(train_subjects, valid_subjects, test_subjects): raise Exception('Set of train, val and test subject must be disjoint.') # use last subject of the non-test subjects for validation train_examples = [] test_examples = [] valid_examples = [] pool = util.BoundedPool(None, 120) if partial_visibility: dir_suffix = '_partial' further_expansion_factor = 1.8 else: dir_suffix = '' if correct_S9 else 'incorrect_S9' further_expansion_factor = 1 for i_subject in [*test_subjects, *train_subjects, *valid_subjects]: if i_subject in train_subjects: examples_container = train_examples elif i_subject in valid_subjects: examples_container = valid_examples else: examples_container = test_examples frame_step = 5 if i_subject in train_subjects else 64 for activity_name, camera_id in itertools.product( get_activity_names(i_subject), range(4)): print(f'Processing S{i_subject} {activity_name} {camera_id}') # Corrupt data in original release: if i_subject == 11 and activity_name == 'Directions' and camera_id == 0: continue data, camera = get_examples(i_subject, activity_name, camera_id, frame_step=frame_step, correct_S9=correct_S9) prev_coords = None for image_relpath, world_coords, bbox in data: # Using very similar examples is wasteful when training. Therefore: # skip frame if all keypoints are within a distance compared to last stored frame. # This is not done when testing, as it would change the results. if (i_subject in train_subjects and prev_coords is not None and np.all( np.linalg.norm(world_coords - prev_coords, axis=1) < 100)): continue prev_coords = world_coords activity_name = activity_name.split(' ')[0] ex = ps3d.Pose3DExample(image_relpath, world_coords, bbox, camera, activity_name=activity_name) pool.apply_async(make_efficient_example, (ex, further_expansion_factor, 1, dir_suffix), callback=examples_container.append) print('Waiting for tasks...') pool.close() pool.join() print('Done...') train_examples.sort(key=lambda x: x.image_path) valid_examples.sort(key=lambda x: x.image_path) test_examples.sort(key=lambda x: x.image_path) return ps3d.Pose3DDataset(joint_info, train_examples, valid_examples, test_examples)
def make_mpii_yolo(): joint_info_full = JointInfo( 'rank,rkne,rhip,lhip,lkne,lank,pelv,thor,neck,head,rwri,relb,rsho,lsho,lelb,lwri', 'lsho-lelb-lwri,rsho-relb-rwri,lhip-lkne-lank,rhip-rkne-rank,neck-head,pelv-thor') joint_info_used = JointInfo( 'rank,rkne,rhip,lhip,lkne,lank,rwri,relb,lelb,lwri', 'lelb-lwri,relb-rwri,lhip-lkne-lank,rhip-rkne-rank') selected_joints = [joint_info_full.ids[name] for name in joint_info_used.names] mat_path = f'{paths.DATA_ROOT}/mpii/mpii_human_pose_v1_u12_1.mat' s = matlabfile.load(mat_path).RELEASE annolist = np.atleast_1d(s.annolist) all_boxes = util.load_pickle(f'{paths.DATA_ROOT}/mpii/yolov3_detections.pkl') examples = [] with util.BoundedPool(None, 120) as pool: for anno_id, (anno, is_train) in enumerate( zip(annolist, util.progressbar(s.img_train))): if not is_train: continue image_path = f'{paths.DATA_ROOT}/mpii/images/{anno.image.name}' annorect = np.atleast_1d(anno.annorect) gt_people = [] for rect_id, rect in enumerate(annorect): if 'annopoints' not in rect or len(rect.annopoints) == 0: continue coords = np.full( shape=[joint_info_full.n_joints, 2], fill_value=np.nan, dtype=np.float32) for joint in np.atleast_1d(rect.annopoints.point): coords[joint.id] = [joint.x, joint.y] bbox = boxlib.expand(boxlib.bb_of_points(coords), 1.25) coords = coords[selected_joints] ex = Pose2DExample(image_path, coords, bbox=bbox) gt_people.append(ex) if not gt_people: continue image_relpath = os.path.relpath(f'images/{anno.image.name}') boxes = [box for box in all_boxes[image_relpath] if box[-1] > 0.5] if not boxes: continue iou_matrix = np.array([[boxlib.iou(gt_person.bbox, box[:4]) for box in boxes] for gt_person in gt_people]) gt_indices, box_indices = scipy.optimize.linear_sum_assignment(-iou_matrix) for i_gt, i_det in zip(gt_indices, box_indices): if iou_matrix[i_gt, i_det] > 0.1: ex = gt_people[i_gt] ex.bbox = np.array(boxes[i_det][:4]) new_im_path = image_path.replace('mpii', 'mpii_downscaled_yolo') without_ext, ext = os.path.splitext(new_im_path) new_im_path = f'{without_ext}_{i_gt:02d}{ext}' pool.apply_async(make_efficient_example, (ex, new_im_path), callback=examples.append) examples.sort(key=lambda ex: ex.image_path) def n_valid_joints(example): return np.count_nonzero(np.all(~np.isnan(example.coords), axis=-1)) examples = [ex for ex in examples if n_valid_joints(ex) > 6] return Pose2DDataset(joint_info_used, examples)
def make_coco(single_person=True): joint_info = JointInfo( 'nose,leye,reye,lear,rear,lsho,rsho,lelb,relb,lwri,rwri,lhip,rhip,lkne,rkne,lank,rank', 'lsho-lelb-lwri,rsho-relb-rwri,lhip-lkne-lank,rhip-rkne-rank,lear-leye-nose-reye-rear') n_joints = joint_info.n_joints learning_phase_shortnames = {TRAIN: 'train', VALID: 'val', TEST: 'test'} UNLABELED = 0 OCCLUDED = 1 VISIBLE = 2 iou_threshold = 0.1 if single_person else 0.5 suffix = '' if single_person else '_multi' examples_per_phase = {TRAIN: [], VALID: []} with util.BoundedPool(None, 120) as pool: for example_phase in (TRAIN, VALID): phase_shortname = learning_phase_shortnames[example_phase] coco_filepath = ( f'{paths.DATA_ROOT}/coco/annotations/person_keypoints_{phase_shortname}2014.json') coco = pycocotools.coco.COCO(coco_filepath) impath_to_examples = {} for ann in coco.anns.values(): filename = coco.imgs[ann['image_id']]['file_name'] image_path = f'{paths.DATA_ROOT}/coco/{phase_shortname}2014/{filename}' joints = np.array(ann['keypoints']).reshape([-1, 3]) visibilities = joints[:, 2] coords = joints[:, :2].astype(np.float32).copy() n_visible_joints = np.count_nonzero(visibilities == VISIBLE) n_occluded_joints = np.count_nonzero(visibilities == OCCLUDED) n_labeled_joints = n_occluded_joints + n_visible_joints if n_visible_joints >= n_joints / 3 and n_labeled_joints >= n_joints / 2: coords[visibilities == UNLABELED] = np.nan bbox_pt1 = np.array(ann['bbox'][0:2], np.float32) bbox_wh = np.array(ann['bbox'][2:4], np.float32) bbox = np.array([*bbox_pt1, *bbox_wh]) ex = Pose2DExample(image_path, coords, bbox=bbox) impath_to_examples.setdefault(image_path, []).append(ex) n_images = len(impath_to_examples) for impath, examples in util.progressbar(impath_to_examples.items(), total=n_images): for i_example, example in enumerate(examples): box = boxlib.expand(boxlib.bb_of_points(example.coords), 1.25) if np.max(box[2:]) < 200: continue if single_person: other_boxes = [boxlib.expand(boxlib.bb_of_points(e.coords), 1.25) for e in examples if e is not example] ious = np.array([boxlib.iou(b, box) for b in other_boxes]) usable = np.all(ious < iou_threshold) else: usable = True if usable: new_im_path = impath.replace('coco', 'coco_downscaled' + suffix) without_ext, ext = os.path.splitext(new_im_path) new_im_path = f'{without_ext}_{i_example:02d}{ext}' pool.apply_async( make_efficient_example, (example, new_im_path), callback=examples_per_phase[example_phase].append) examples_per_phase[TRAIN].sort(key=lambda ex: ex.image_path) examples_per_phase[VALID].sort(key=lambda ex: ex.image_path) return Pose2DDataset(joint_info, examples_per_phase[TRAIN], examples_per_phase[VALID])
def generate_all_overall_masks(composites, i_valids): with util.BoundedPool(None, 120) as pool: for i, sample in zip(i_valids, util.progressbar(composites[i_valids])): pool.apply_async(save_overall_mask, (sample, i))