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))
