def test_imshow_keypoints():
# 2D keypoint
img = np.zeros((100, 100, 3), dtype=np.uint8)
kpts = np.array([[1, 1, 1], [10, 10, 1]], dtype=np.float32)
pose_result = [kpts]
skeleton = [[0, 1]]
pose_kpt_color = [(127, 127, 127)] * len(kpts)
pose_link_color = [(127, 127, 127)] * len(skeleton)
img_vis_2d = imshow_keypoints(
img,
pose_result,
skeleton=skeleton,
pose_kpt_color=pose_kpt_color,
pose_link_color=pose_link_color,
show_keypoint_weight=True)
# 3D keypoint
kpts_3d = np.array([[0, 0, 0, 1], [1, 1, 1, 1]], dtype=np.float32)
pose_result_3d = [{'keypoints_3d': kpts_3d, 'title': 'test'}]
_ = imshow_keypoints_3d(
pose_result_3d,
img=img_vis_2d,
skeleton=skeleton,
pose_kpt_color=pose_kpt_color,
pose_link_color=pose_link_color,
vis_height=400)
def _draw_keypoint(self, canvas: np.ndarray, input_msg: FrameMessage):
"""Draw object keypoints."""
objects = input_msg.get_objects(lambda x: 'pose_model_cfg' in x)
# return if there is no object with keypoints
if not objects:
return canvas
for model_cfg, group in groupby(objects,
lambda x: x['pose_model_cfg']):
dataset_info = DatasetInfo(model_cfg.dataset_info)
keypoints = [obj['keypoints'] for obj in group]
imshow_keypoints(canvas,
keypoints,
skeleton=dataset_info.skeleton,
kpt_score_thr=0.3,
pose_kpt_color=dataset_info.pose_kpt_color,
pose_link_color=dataset_info.pose_link_color,
radius=self.radius,
thickness=self.thickness)
return canvas
def test_imshow_keypoints_2d():
img = np.zeros((100, 100, 3), dtype=np.uint8)
kpts = np.array([[1, 1, 1], [2, 2, 1], [4, 4, 1], [8, 8, 1]],
dtype=np.float32)
pose_result = [kpts]
skeleton = [[0, 1], [1, 2], [2, 3]]
# None: kpt or link is hidden
pose_kpt_color = [None] + [(127, 127, 127)] * (len(kpts) - 1)
pose_link_color = [(127, 127, 127)] * (len(skeleton) - 1) + [None]
_ = imshow_keypoints(img,
pose_result,
skeleton=skeleton,
pose_kpt_color=pose_kpt_color,
pose_link_color=pose_link_color,
show_keypoint_weight=True)
def show_result(self,
img,
result,
skeleton=None,
kpt_score_thr=0.3,
bbox_color='green',
pose_kpt_color=None,
pose_link_color=None,
text_color='white',
radius=4,
thickness=1,
font_scale=0.5,
bbox_thickness=1,
win_name='',
show=False,
show_keypoint_weight=False,
wait_time=0,
out_file=None):
"""Draw `result` over `img`.
Args:
img (str or Tensor): The image to be displayed.
result (list[dict]): The results to draw over `img`
(bbox_result, pose_result).
skeleton (list[list]): The connection of keypoints.
skeleton is 0-based indexing.
kpt_score_thr (float, optional): Minimum score of keypoints
to be shown. Default: 0.3.
bbox_color (str or tuple or :obj:`Color`): Color of bbox lines.
pose_kpt_color (np.array[Nx3]`): Color of N keypoints.
If None, do not draw keypoints.
pose_link_color (np.array[Mx3]): Color of M links.
If None, do not draw links.
text_color (str or tuple or :obj:`Color`): Color of texts.
radius (int): Radius of circles.
thickness (int): Thickness of lines.
font_scale (float): Font scales of texts.
win_name (str): The window name.
show (bool): Whether to show the image. Default: False.
show_keypoint_weight (bool): Whether to change the transparency
using the predicted confidence scores of keypoints.
wait_time (int): Value of waitKey param.
Default: 0.
out_file (str or None): The filename to write the image.
Default: None.
Returns:
Tensor: Visualized img, only if not `show` or `out_file`.
"""
img = mmcv.imread(img)
img = img.copy()
bbox_result = []
bbox_labels = []
pose_result = []
for res in result:
if 'bbox' in res:
bbox_result.append(res['bbox'])
bbox_labels.append(res.get('label', None))
pose_result.append(res['keypoints'])
if bbox_result:
bboxes = np.vstack(bbox_result)
# draw bounding boxes
imshow_bboxes(img,
bboxes,
labels=bbox_labels,
colors=bbox_color,
text_color=text_color,
thickness=bbox_thickness,
font_scale=font_scale,
show=False)
if pose_result:
imshow_keypoints(img, pose_result, skeleton, kpt_score_thr,
pose_kpt_color, pose_link_color, radius,
thickness)
if show:
imshow(img, win_name, wait_time)
if out_file is not None:
imwrite(img, out_file)
return img
def show_result(self,
result,
img=None,
skeleton=None,
pose_kpt_color=None,
pose_link_color=None,
radius=8,
thickness=2,
vis_height=400,
num_instances=-1,
axis_azimuth=70,
win_name='',
show=False,
wait_time=0,
out_file=None):
"""Visualize 3D pose estimation results.
Args:
result (list[dict]): The pose estimation results containing:
- "keypoints_3d" ([K,4]): 3D keypoints
- "keypoints" ([K,3] or [T,K,3]): Optional for visualizing
2D inputs. If a sequence is given, only the last frame
will be used for visualization
- "bbox" ([4,] or [T,4]): Optional for visualizing 2D inputs
- "title" (str): title for the subplot
img (str or Tensor): Optional. The image to visualize 2D inputs on.
skeleton (list of [idx_i,idx_j]): Skeleton described by a list of
links, each is a pair of joint indices.
pose_kpt_color (np.array[Nx3]`): Color of N keypoints.
If None, do not draw keypoints.
pose_link_color (np.array[Mx3]): Color of M links.
If None, do not draw links.
radius (int): Radius of circles.
thickness (int): Thickness of lines.
vis_height (int): The image height of the visualization. The width
will be N*vis_height depending on the number of visualized
items.
num_instances (int): Number of instances to be shown in 3D. If
smaller than 0, all the instances in the result will be shown.
Otherwise, pad or truncate the result to a length of
num_instances.
axis_azimuth (float): axis azimuth angle for 3D visualizations.
win_name (str): The window name.
show (bool): Whether to directly show the visualization.
wait_time (int): Value of waitKey param.
Default: 0.
out_file (str or None): The filename to write the image.
Default: None.
Returns:
Tensor: Visualized img, only if not `show` or `out_file`.
"""
if num_instances < 0:
assert len(result) > 0
result = sorted(result, key=lambda x: x.get('track_id', 1e4))
# draw image and input 2d poses
if img is not None:
img = mmcv.imread(img)
bbox_result = []
pose_input_2d = []
for res in result:
if 'bbox' in res:
bbox = np.array(res['bbox'])
if bbox.ndim != 1:
assert bbox.ndim == 2
bbox = bbox[-1] # Get bbox from the last frame
bbox_result.append(bbox)
if 'keypoints' in res:
kpts = np.array(res['keypoints'])
if kpts.ndim != 2:
assert kpts.ndim == 3
kpts = kpts[-1] # Get 2D keypoints from the last frame
pose_input_2d.append(kpts)
if len(bbox_result) > 0:
bboxes = np.vstack(bbox_result)
imshow_bboxes(
img,
bboxes,
colors='green',
thickness=thickness,
show=False)
if len(pose_input_2d) > 0:
imshow_keypoints(
img,
pose_input_2d,
skeleton,
kpt_score_thr=0.3,
pose_kpt_color=pose_kpt_color,
pose_link_color=pose_link_color,
radius=radius,
thickness=thickness)
img = mmcv.imrescale(img, scale=vis_height / img.shape[0])
img_vis = imshow_keypoints_3d(
result,
img,
skeleton,
pose_kpt_color,
pose_link_color,
vis_height,
num_instances=num_instances,
axis_azimuth=axis_azimuth,
)
if show:
mmcv.visualization.imshow(img_vis, win_name, wait_time)
if out_file is not None:
mmcv.imwrite(img_vis, out_file)
return img_vis
def show_result(self,
img,
img_metas,
visualize_2d=False,
input_heatmaps=None,
dataset_info=None,
radius=4,
thickness=2,
out_dir=None,
show=False):
"""Visualize the results."""
result = self.forward_test(img,
img_metas,
input_heatmaps=input_heatmaps)
pose_3d = result['pose_3d']
sample_id = result['sample_id']
batch_size = pose_3d.shape[0]
# get kpts and skeleton structure
for i in range(batch_size):
# visualize 3d results
img_meta = img_metas[i]
num_cameras = len(img_meta['camera'])
pose_3d_i = pose_3d[i]
pose_3d_i = pose_3d_i[pose_3d_i[:, 0, 3] >= 0]
num_persons, num_keypoints, _ = pose_3d_i.shape
pose_3d_list = [p[..., [0, 1, 2, 4]]
for p in pose_3d_i] if num_persons > 0 else []
img_3d = imshow_multiview_keypoints_3d(
pose_3d_list,
skeleton=dataset_info.skeleton,
pose_kpt_color=dataset_info.pose_kpt_color[:num_keypoints],
pose_link_color=dataset_info.pose_link_color,
space_size=self.human_detector.space_size,
space_center=self.human_detector.space_center)
if out_dir is not None:
mmcv.image.imwrite(
img_3d,
os.path.join(out_dir, 'vis_3d', f'{sample_id[i]}_3d.jpg'))
if visualize_2d:
for j in range(num_cameras):
single_camera = SimpleCamera(img_meta['camera'][j])
# img = mmcv.imread(img)
if num_persons > 0:
pose_2d = np.ones_like(pose_3d_i[..., :3])
pose_2d_flat = single_camera.world_to_pixel(
pose_3d_i[..., :3].reshape((-1, 3)))
pose_2d[..., :2] = pose_2d_flat.reshape(
(num_persons, -1, 2))
pose_2d_list = [pose for pose in pose_2d]
else:
pose_2d_list = []
with tempfile.TemporaryDirectory() as tmpdir:
if 'image_file' in img_meta:
img_file = img_meta['image_file'][j]
else:
img_size = img_meta['center'][j] * 2
img = np.zeros(
[int(img_size[1]),
int(img_size[0]), 3],
dtype=np.uint8)
img.fill(255) # or img[:] = 255
img_file = os.path.join(tmpdir, 'tmp.jpg')
mmcv.image.imwrite(img, img_file)
img = imshow_keypoints(
img_file, pose_2d_list, dataset_info.skeleton, 0.0,
dataset_info.pose_kpt_color[:num_keypoints],
dataset_info.pose_link_color, radius, thickness)
if out_dir is not None:
mmcv.image.imwrite(
img,
os.path.join(out_dir,
f'{sample_id[i]}_{j}_2d.jpg'))
