def get_chair_mask(i_subject, i_seq, i_cam, i_frame):
p = f'{paths.DATA_ROOT}/3dhp/S{i_subject}/Seq{i_seq}/ChairMasks/img_{i_cam}_{i_frame:06d}.jpg'
chroma_frame = improc.imread_jpeg(p)
is_fg = chroma_frame[..., 0] < 32
mask, objbox = improc.largest_connected_component(is_fg)
encoded_mask = improc.encode_mask(mask)
return encoded_mask
def get_mask(i_subject, i_seq, i_cam, i_frame):
chroma_frame = improc.imread_jpeg(
f'{paths.DATA_ROOT}/3dhp/S{i_subject}/Seq{i_seq}/FGmasks/img_{i_cam}_{i_frame:06d}.jpg'
)
person_box = get_box(i_subject, i_seq, i_cam, i_frame)
is_fg = chroma_frame[..., 0] > 100
n_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
is_fg.astype(np.uint8), 4, cv2.CV_32S)
component_boxes = stats[:, :4]
ious = [
boxlib.iou(component_box, person_box)
for component_box in component_boxes
]
ious[0] = 0
person_label = np.argmax(ious)
mask = (labels == person_label).astype(np.uint8)
# Remove foreground pixels that are far from the person box
intbox = boxlib.intersect(boxlib.full_box((2048, 2048)),
boxlib.expand(person_box, 1.3)).astype(int)
mask[:intbox[1]] = 0
mask[:, :intbox[0]] = 0
mask[:, intbox[0] + intbox[2]:] = 0
mask[intbox[1] + intbox[3]:] = 0
encoded_mask = improc.encode_mask(mask)
return encoded_mask
def get_connected_component_with_highest_iou(mask, person_box):
"""Finds the 4-connected component in `mask` with the highest bbox IoU with the `person box`"""
mask = mask.astype(np.uint8)
_, labels, stats, _ = cv2.connectedComponentsWithStats(mask, 4, cv2.CV_32S)
component_boxes = stats[:, :4]
ious = [boxlib.iou(component_box, person_box) for component_box in component_boxes]
person_label = np.argmax(ious)
return improc.encode_mask(labels == person_label)
def make_efficient_example(ex, root_muco, i_person):
image_relpath = ex.image_path
max_rotate = np.pi / 6
padding_factor = 1 / 0.85
scale_up_factor = 1 / 0.85
scale_down_factor = 1 / 0.85
shift_factor = 1.2
base_dst_side = 256
box_center = boxlib.center(ex.bbox)
s = np.sin(max_rotate)
c = np.cos(max_rotate)
rot_bbox_size = (np.array([[c, s], [s, c]]) @ ex.bbox[2:, np.newaxis])[:, 0]
side = np.max(rot_bbox_size)
rot_bbox_size = np.array([side, side])
rot_bbox = boxlib.box_around(box_center, rot_bbox_size)
scale_factor = min(base_dst_side / np.max(ex.bbox[2:]) * scale_up_factor, 1)
expansion_factor = padding_factor * shift_factor * scale_down_factor
expanded_bbox = boxlib.expand(rot_bbox, expansion_factor)
expanded_bbox = boxlib.intersect(expanded_bbox, boxlib.full_box([2048, 2048]))
new_camera = ex.camera.copy()
new_camera.intrinsic_matrix[:2, 2] -= expanded_bbox[:2]
new_camera.scale_output(scale_factor)
new_camera.undistort()
dst_shape = improc.rounded_int_tuple(scale_factor * expanded_bbox[[3, 2]])
new_im_path = f'{root_muco}_downscaled/{image_relpath[:-4]}_{i_person:01d}.jpg'
if not (util.is_file_newer(new_im_path, "2020-02-15T23:28:26")):
im = improc.imread_jpeg(f'{root_muco}/{image_relpath}')
new_im = cameralib.reproject_image(im, ex.camera, new_camera, dst_shape, antialias_factor=4)
util.ensure_path_exists(new_im_path)
imageio.imwrite(new_im_path, new_im, quality=95)
new_bbox_topleft = cameralib.reproject_image_points(ex.bbox[:2], ex.camera, new_camera)
new_bbox = np.concatenate([new_bbox_topleft, ex.bbox[2:] * scale_factor])
if ex.mask is None:
noext, ext = os.path.splitext(image_relpath[:-4])
noext = noext.replace('unaugmented_set_001/', '')
mask = improc.decode_mask(util.load_pickle(f'{root_muco}/masks/{noext}.pkl'))
else:
mask = ex.mask
if mask is False:
new_mask_encoded = None
else:
new_mask = cameralib.reproject_image(mask, ex.camera, new_camera, dst_shape)
new_mask_encoded = improc.encode_mask(new_mask)
return p3ds.Pose3DExample(
os.path.relpath(new_im_path, paths.DATA_ROOT), ex.world_coords.astype(np.float32),
new_bbox.astype(np.float32), new_camera, mask=new_mask_encoded,
univ_coords=ex.univ_coords.astype(np.float32))
def save_overall_mask(people, i_out):
def _get_mask(i):
i_subject, i_seq, i_cam, i_frame = people[i]
encoded_mask = get_mask(int(i_subject), int(i_seq), int(i_cam),
int(i_frame))
return improc.decode_mask(encoded_mask)
def _get_chair_mask(i):
i_subject, i_seq, i_cam, i_frame = people[i]
encoded_mask = get_chair_mask(int(i_subject), int(i_seq), int(i_cam),
int(i_frame))
return improc.decode_mask(encoded_mask)
overall_mask = _get_mask(0)
np.maximum(overall_mask, _get_chair_mask(0), out=overall_mask)
for i in range(1, 4):
np.maximum(overall_mask, _get_mask(i), out=overall_mask)
np.maximum(overall_mask, _get_chair_mask(i), out=overall_mask)
s = f'{i_out + 1:06d}'
out_path = f'{paths.DATA_ROOT}/muco/masks/{s[:2]}/{s[:4]}/{s}.pkl'
util.dump_pickle(improc.encode_mask(overall_mask), out_path)
def empty_mask():
return improc.encode_mask(np.zeros((2048, 2048), dtype=np.uint8))