def gen_bottom_up_color_skeleton(proposals,
bu_hor,
bu_ver,
gt_pose,
file_dir,
base_time,
verbose=True):
ensure_dir(file_dir)
image_tag = "{h_v:}_{time:05d}_{gt:}"
bu_hor, bu_ver = bu_hor[0], bu_ver[0]
def _yield_record():
for t in range(bu_hor.shape[1]):
yield t, bu_hor[:, t], bu_ver[:, t], gt_pose[t]
map_in_pool(
fn=partial(gen_bottom_up_color_skeleton_worker,
image_tag=image_tag,
file_dir=file_dir,
base_time=base_time,
verbose=verbose,
color=KPTS_COLOR[:, np.newaxis, np.newaxis, :]),
data=_yield_record(),
single_process=False,
verbose=verbose,
nr_pro=8,
)
def gen_bottom_up_heatmap(proposals,
bu_hor,
bu_ver,
gt_pose,
file_dir,
base_time,
verbose=True):
ensure_dir(file_dir)
image_tag = "{kpt_name:}_{h_v:}_{time:05d}"
bu_hor, bu_ver = bu_hor[0], bu_ver[0]
def yield_record():
for t in tqdm(range(bu_hor.shape[1])):
for i_kpt, kpt_name in enumerate(pose_config.KPT_NAMES):
yield t, i_kpt, kpt_name, bu_hor[i_kpt][t], bu_ver[i_kpt][t]
map_in_pool(fn=partial(gen_bottom_up_heatmap_worker,
image_tag=image_tag,
file_dir=file_dir,
base_time=base_time,
verbose=True),
data=yield_record(),
single_process=False,
verbose=verbose,
nr_pro=8)
def gen_3d_pose_vis_image(pred_pose,
gt_pose,
file_dir,
base_time,
verbose=True):
"""
Generate 3d pose visualization images for a time window.
:param pred_pose:
:param gt_pose:
:param heatmap:
:param proposal:
:param filename:
:return:
"""
ensure_dir(file_dir)
time_range = len(pred_pose)
def yield_record():
for t in tqdm(range(time_range)):
yield pred_pose[t], gt_pose[t], os.path.join(
file_dir, "gt_pred_{:05d}".format(base_time + t)), t
map_in_pool(fn=partial(_worker_draw_poses_with_gt, verbose=verbose),
data=yield_record(),
single_process=False,
nr_pro=12,
verbose=verbose)
gc.collect()
def vis_3d_heatmap(heatmap, file_path):
assert len(heatmap.shape) == 3
X, Y, Z = heatmap.shape
ensure_dir(file_path)
# use multi process when heat map is large.
if max(Y, X) > 16:
_ = map_in_pool(
fn=partial(vis_3d_heatmap_worker_y, verbose=True, file_path=file_path),
data=vis_3d_heatmap_distributor_y(heatmap, Y),
single_process=False,
verbose=True,
nr_pro=8
)
_ = map_in_pool(
fn=partial(vis_3d_heatmap_worker_x, verbose=True, file_path=file_path),
data=vis_3d_heatmap_distributor_x(heatmap, X),
single_process=False,
verbose=True,
nr_pro=8
)
else:
# visualize horizontal plane
for y in range(Y):
vis_2d_heatmap(heatmap=heatmap[:, y, :], filename=os.path.join(file_path, "y_{}".format(y)))
# visualize vertical plane
for x in range(X):
vis_2d_heatmap(heatmap=heatmap[x, :, :], filename=os.path.join(file_path, "x_{}".format(x)))
def gen_aligned_feature_map(proposals,
feature_hor,
feature_ver,
file_dir,
base_time,
verbose=True):
ensure_dir(file_dir)
image_tag = "{h_v:}_{time:05d}"
feature_hor, feature_ver = np.moveaxis(feature_hor[0], 0, 1), np.moveaxis(
feature_ver[0], 0, 1)
def yield_records():
for t in range(len(proposals)):
yield t, proposals[t], feature_hor[t], feature_ver[t]
_ = map_in_pool(fn=partial(gen_aligned_feature_map_worker,
file_dir=file_dir,
base_time=base_time,
image_tag=image_tag,
verbose=verbose),
data=yield_records(),
single_process=False,
verbose=verbose,
nr_pro=16)
def gen_aligned_pose_heatmap(proposals,
heatmap_h,
heatmap_v,
gt_pose,
pred_pose,
file_dir,
base_time,
verbose=True):
"""
plot and save aligned pose heatmap.
:param proposals: list (t, n_prop, 4)
:param heatmap_h: list (t, n_prop, (14, 48, 48) )
:param heatmap_v: list (t, n_prop, (14, 128, 48) )
:return:
"""
ensure_dir(file_dir)
image_tag = "{kpt_name:}_{h_v:}_{time:05d}.png"
def yeild_records():
r_length = len(proposals)
for t in range(r_length):
# align gt pose to pred pose
gt_threshold = 10
gt_pose_candidates, gt_poses_aligned = gt_pose[t], []
for pose in pred_pose[t]:
# for each pred pose, find nearest one.
dist = np.asarray([
pose_dist(pose, p) for p in gt_pose_candidates
if np.sum(p) > 0
])
dist_min_idx = np.argmin(dist)
# handle no gt with threshold.
if dist[dist_min_idx] > gt_threshold:
gt_poses_aligned.append(
np.zeros_like(gt_pose_candidates[dist_min_idx]))
else:
gt_poses_aligned.append(
np.array(gt_pose_candidates[dist_min_idx]))
# distribute work for visualization.
for idx_kpt in range(pose_config.N_KPTS):
# number of prop
nr_prop = len(proposals[t])
# proposals, heatmap, keypoint
proposals_t, heatmap_h_t, heatmap_v_t = proposals[
t], heatmap_h[t], heatmap_v[t]
kpt_pred_list = [
pred_pose[t][i][idx_kpt] for i in range(nr_prop)
]
kpt_gt_list = [
gt_poses_aligned[i][idx_kpt] for i in range(nr_prop)
]
yield (t, idx_kpt, proposals_t, heatmap_h_t, heatmap_v_t,
kpt_pred_list, kpt_gt_list)
_ = map_in_pool(fn=partial(_gen_pose_heatmap_per_t_kpt,
file_dir=file_dir,
base_time=base_time,
image_tag=image_tag,
verbose=verbose),
data=yeild_records(),
single_process=False,
verbose=verbose,
nr_pro=16)
gc.collect()