def test_imshow_bbox():
img = np.zeros((100, 100, 3), dtype=np.uint8)
bboxes = np.array([[10, 10, 30, 30], [10, 50, 30, 80]], dtype=np.float32)
labels = ['label 1', 'label 2']
colors = ['red', 'green']
with tempfile.TemporaryDirectory() as tmpdir:
_ = imshow_bboxes(img,
bboxes,
labels=labels,
colors=colors,
show=False,
out_file=f'{tmpdir}/out.png')
# test case of empty bboxes
_ = imshow_bboxes(img,
np.zeros((0, 4), dtype=np.float32),
labels=None,
colors='red',
show=False)
# test unmatched bboxes and labels
with pytest.raises(AssertionError):
_ = imshow_bboxes(img,
np.zeros((0, 4), dtype=np.float32),
labels=labels[:1],
colors='red',
show=False)
def _draw_bbox(self, canvas: np.ndarray, input_msg: FrameMessage):
"""Draw object bboxes."""
if self.must_have_keypoint:
objects = input_msg.get_objects(
lambda x: 'bbox' in x and 'keypoints' in x)
else:
objects = input_msg.get_objects(lambda x: 'bbox' in x)
# return if there is no detected objects
if not objects:
return canvas
bboxes = [obj['bbox'] for obj in objects]
labels = [obj.get('label', None) for obj in objects]
default_color = (0, 255, 0)
# Get bbox colors
if isinstance(self.bbox_color, dict):
colors = [
self.bbox_color.get(label, default_color) for label in labels
]
else:
colors = self.bbox_color
imshow_bboxes(canvas,
np.vstack(bboxes),
labels=labels,
colors=colors,
text_color='white',
font_scale=0.5,
show=False)
return canvas
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