# Copyright (c) OpenMMLab. All rights reserved.
import math
from typing import Dict, List, Optional, Tuple, Union
import cv2
import mmcv
import numpy as np
import torch
from mmengine.dist import master_only
from mmengine.structures import InstanceData, PixelData
from mmpose.datasets.datasets.utils import parse_pose_metainfo
from mmpose.registry import VISUALIZERS
from mmpose.structures import PoseDataSample
from .opencv_backend_visualizer import OpencvBackendVisualizer
from .simcc_vis import SimCCVisualizer
def _get_adaptive_scales(areas: np.ndarray,
min_area: int = 800,
max_area: int = 30000) -> np.ndarray:
"""Get adaptive scales according to areas.
The scale range is [0.5, 1.0]. When the area is less than
``min_area``, the scale is 0.5 while the area is larger than
``max_area``, the scale is 1.0.
Args:
areas (ndarray): The areas of bboxes or masks with the
shape of (n, ).
min_area (int): Lower bound areas for adaptive scales.
Defaults to 800.
max_area (int): Upper bound areas for adaptive scales.
Defaults to 30000.
Returns:
ndarray: The adaotive scales with the shape of (n, ).
"""
scales = 0.5 + (areas - min_area) / (max_area - min_area)
scales = np.clip(scales, 0.5, 1.0)
return scales
@VISUALIZERS.register_module()
class PoseLocalVisualizer(OpencvBackendVisualizer):
"""MMPose Local Visualizer.
Args:
name (str): Name of the instance. Defaults to 'visualizer'.
image (np.ndarray, optional): the origin image to draw. The format
should be RGB. Defaults to ``None``
vis_backends (list, optional): Visual backend config list. Defaults to
``None``
save_dir (str, optional): Save file dir for all storage backends.
If it is ``None``, the backend storage will not save any data.
Defaults to ``None``
bbox_color (str, tuple(int), optional): Color of bbox lines.
The tuple of color should be in BGR order. Defaults to ``'green'``
kpt_color (str, tuple(tuple(int)), optional): Color of keypoints.
The tuple of color should be in BGR order. Defaults to ``'red'``
link_color (str, tuple(tuple(int)), optional): Color of skeleton.
The tuple of color should be in BGR order. Defaults to ``None``
line_width (int, float): The width of lines. Defaults to 1
radius (int, float): The radius of keypoints. Defaults to 4
show_keypoint_weight (bool): Whether to adjust the transparency
of keypoints according to their score. Defaults to ``False``
alpha (int, float): The transparency of bboxes. Defaults to ``0.8``
Examples:
>>> import numpy as np
>>> from mmengine.structures import InstanceData
>>> from mmpose.structures import PoseDataSample
>>> from mmpose.visualization import PoseLocalVisualizer
>>> pose_local_visualizer = PoseLocalVisualizer(radius=1)
>>> image = np.random.randint(0, 256,
... size=(10, 12, 3)).astype('uint8')
>>> gt_instances = InstanceData()
>>> gt_instances.keypoints = np.array([[[1, 1], [2, 2], [4, 4],
... [8, 8]]])
>>> gt_pose_data_sample = PoseDataSample()
>>> gt_pose_data_sample.gt_instances = gt_instances
>>> dataset_meta = {'skeleton_links': [[0, 1], [1, 2], [2, 3]]}
>>> pose_local_visualizer.set_dataset_meta(dataset_meta)
>>> pose_local_visualizer.add_datasample('image', image,
... gt_pose_data_sample)
>>> pose_local_visualizer.add_datasample(
... 'image', image, gt_pose_data_sample,
... out_file='out_file.jpg')
>>> pose_local_visualizer.add_datasample(
... 'image', image, gt_pose_data_sample,
... show=True)
>>> pred_instances = InstanceData()
>>> pred_instances.keypoints = np.array([[[1, 1], [2, 2], [4, 4],
... [8, 8]]])
>>> pred_instances.score = np.array([0.8, 1, 0.9, 1])
>>> pred_pose_data_sample = PoseDataSample()
>>> pred_pose_data_sample.pred_instances = pred_instances
>>> pose_local_visualizer.add_datasample('image', image,
... gt_pose_data_sample,
... pred_pose_data_sample)
"""
def __init__(self,
name: str = 'visualizer',
image: Optional[np.ndarray] = None,
vis_backends: Optional[Dict] = None,
save_dir: Optional[str] = None,
bbox_color: Optional[Union[str, Tuple[int]]] = 'green',
kpt_color: Optional[Union[str, Tuple[Tuple[int]]]] = 'red',
link_color: Optional[Union[str, Tuple[Tuple[int]]]] = None,
text_color: Optional[Union[str,
Tuple[int]]] = (255, 255, 255),
skeleton: Optional[Union[List, Tuple]] = None,
line_width: Union[int, float] = 1,
radius: Union[int, float] = 3,
show_keypoint_weight: bool = False,
backend: str = 'opencv',
alpha: float = 0.8):
super().__init__(
name=name,
image=image,
vis_backends=vis_backends,
save_dir=save_dir,
backend=backend)
self.bbox_color = bbox_color
self.kpt_color = kpt_color
self.link_color = link_color
self.line_width = line_width
self.text_color = text_color
self.skeleton = skeleton
self.radius = radius
self.alpha = alpha
self.show_keypoint_weight = show_keypoint_weight
# Set default value. When calling
# `PoseLocalVisualizer().set_dataset_meta(xxx)`,
# it will override the default value.
self.dataset_meta = {}
def set_dataset_meta(self,
dataset_meta: Dict,
skeleton_style: str = 'mmpose'):
"""Assign dataset_meta to the visualizer. The default visualization
settings will be overridden.
Args:
dataset_meta (dict): meta information of dataset.
"""
if dataset_meta.get(
'dataset_name') == 'coco' and skeleton_style == 'openpose':
dataset_meta = parse_pose_metainfo(
dict(from_file='configs/_base_/datasets/coco_openpose.py'))
if isinstance(dataset_meta, dict):
self.dataset_meta = dataset_meta.copy()
self.bbox_color = dataset_meta.get('bbox_color', self.bbox_color)
self.kpt_color = dataset_meta.get('keypoint_colors',
self.kpt_color)
self.link_color = dataset_meta.get('skeleton_link_colors',
self.link_color)
self.skeleton = dataset_meta.get('skeleton_links', self.skeleton)
# sometimes self.dataset_meta is manually set, which might be None.
# it should be converted to a dict at these times
if self.dataset_meta is None:
self.dataset_meta = {}
def _draw_instances_bbox(self, image: np.ndarray,
instances: InstanceData) -> np.ndarray:
"""Draw bounding boxes and corresponding labels of GT or prediction.
Args:
image (np.ndarray): The image to draw.
instances (:obj:`InstanceData`): Data structure for
instance-level annotations or predictions.
Returns:
np.ndarray: the drawn image which channel is RGB.
"""
self.set_image(image)
if 'bboxes' in instances:
bboxes = instances.bboxes
self.draw_bboxes(
bboxes,
edge_colors=self.bbox_color,
alpha=self.alpha,
line_widths=self.line_width)
else:
return self.get_image()
if 'labels' in instances and self.text_color is not None:
classes = self.dataset_meta.get('classes', None)
labels = instances.labels
positions = bboxes[:, :2]
areas = (bboxes[:, 3] - bboxes[:, 1]) * (
bboxes[:, 2] - bboxes[:, 0])
scales = _get_adaptive_scales(areas)
for i, (pos, label) in enumerate(zip(positions, labels)):
label_text = classes[
label] if classes is not None else f'class {label}'
if isinstance(self.bbox_color,
tuple) and max(self.bbox_color) > 1:
facecolor = [c / 255.0 for c in self.bbox_color]
else:
facecolor = self.bbox_color
self.draw_texts(
label_text,
pos,
colors=self.text_color,
font_sizes=int(13 * scales[i]),
vertical_alignments='bottom',
bboxes=[{
'facecolor': facecolor,
'alpha': 0.8,
'pad': 0.7,
'edgecolor': 'none'
}])
return self.get_image()
def _draw_instances_kpts(self,
image: np.ndarray,
instances: InstanceData,
kpt_thr: float = 0.3,
show_kpt_idx: bool = False,
skeleton_style: str = 'mmpose'):
"""Draw keypoints and skeletons (optional) of GT or prediction.
Args:
image (np.ndarray): The image to draw.
instances (:obj:`InstanceData`): Data structure for
instance-level annotations or predictions.
kpt_thr (float, optional): Minimum threshold of keypoints
to be shown. Default: 0.3.
show_kpt_idx (bool): Whether to show the index of keypoints.
Defaults to ``False``
skeleton_style (str): Skeleton style selection. Defaults to
``'mmpose'``
Returns:
np.ndarray: the drawn image which channel is RGB.
"""
self.set_image(image)
img_h, img_w, _ = image.shape
if 'keypoints' in instances:
keypoints = instances.get('transformed_keypoints',
instances.keypoints)
if 'keypoint_scores' in instances:
scores = instances.keypoint_scores
else:
scores = np.ones(keypoints.shape[:-1])
if 'keypoints_visible' in instances:
keypoints_visible = instances.keypoints_visible
else:
keypoints_visible = np.ones(keypoints.shape[:-1])
if skeleton_style == 'openpose':
keypoints_info = np.concatenate(
(keypoints, scores[..., None], keypoints_visible[...,
None]),
axis=-1)
# compute neck joint
neck = np.mean(keypoints_info[:, [5, 6]], axis=1)
# neck score when visualizing pred
neck[:, 2:4] = np.logical_and(
keypoints_info[:, 5, 2:4] > kpt_thr,
keypoints_info[:, 6, 2:4] > kpt_thr).astype(int)
new_keypoints_info = np.insert(
keypoints_info, 17, neck, axis=1)
mmpose_idx = [
17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3
]
openpose_idx = [
1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17
]
new_keypoints_info[:, openpose_idx] = \
new_keypoints_info[:, mmpose_idx]
keypoints_info = new_keypoints_info
keypoints, scores, keypoints_visible = keypoints_info[
..., :2], keypoints_info[..., 2], keypoints_info[..., 3]
for kpts, score, visible in zip(keypoints, scores,
keypoints_visible):
kpts = np.array(kpts, copy=False)
if self.kpt_color is None or isinstance(self.kpt_color, str):
kpt_color = [self.kpt_color] * len(kpts)
elif len(self.kpt_color) == len(kpts):
kpt_color = self.kpt_color
else:
raise ValueError(
f'the length of kpt_color '
f'({len(self.kpt_color)}) does not matches '
f'that of keypoints ({len(kpts)})')
# draw links
if self.skeleton is not None and self.link_color is not None:
if self.link_color is None or isinstance(
self.link_color, str):
link_color = [self.link_color] * len(self.skeleton)
elif len(self.link_color) == len(self.skeleton):
link_color = self.link_color
else:
raise ValueError(
f'the length of link_color '
f'({len(self.link_color)}) does not matches '
f'that of skeleton ({len(self.skeleton)})')
for sk_id, sk in enumerate(self.skeleton):
pos1 = (int(kpts[sk[0], 0]), int(kpts[sk[0], 1]))
pos2 = (int(kpts[sk[1], 0]), int(kpts[sk[1], 1]))
if not (visible[sk[0]] and visible[sk[1]]):
continue
if (pos1[0] <= 0 or pos1[0] >= img_w or pos1[1] <= 0
or pos1[1] >= img_h or pos2[0] <= 0
or pos2[0] >= img_w or pos2[1] <= 0
or pos2[1] >= img_h or score[sk[0]] < kpt_thr
or score[sk[1]] < kpt_thr
or link_color[sk_id] is None):
# skip the link that should not be drawn
continue
X = np.array((pos1[0], pos2[0]))
Y = np.array((pos1[1], pos2[1]))
color = link_color[sk_id]
if not isinstance(color, str):
color = tuple(int(c) for c in color)
transparency = self.alpha
if self.show_keypoint_weight:
transparency *= max(
0, min(1, 0.5 * (score[sk[0]] + score[sk[1]])))
if skeleton_style == 'openpose':
mX = np.mean(X)
mY = np.mean(Y)
length = ((Y[0] - Y[1])**2 + (X[0] - X[1])**2)**0.5
angle = math.degrees(
math.atan2(Y[0] - Y[1], X[0] - X[1]))
stickwidth = 2
polygons = cv2.ellipse2Poly(
(int(mX), int(mY)),
(int(length / 2), int(stickwidth)), int(angle),
0, 360, 1)
self.draw_polygons(
polygons,
edge_colors=color,
face_colors=color,
alpha=transparency)
else:
self.draw_lines(
X, Y, color, line_widths=self.line_width)
# draw each point on image
for kid, kpt in enumerate(kpts):
if score[kid] < kpt_thr or not visible[
kid] or kpt_color[kid] is None:
# skip the point that should not be drawn
continue
color = kpt_color[kid]
if not isinstance(color, str):
color = tuple(int(c) for c in color)
transparency = self.alpha
if self.show_keypoint_weight:
transparency *= max(0, min(1, score[kid]))
self.draw_circles(
kpt,
radius=np.array([self.radius]),
face_colors=color,
edge_colors=color,
alpha=transparency,
line_widths=self.radius)
if show_kpt_idx:
kpt[0] += self.radius
kpt[1] -= self.radius
self.draw_texts(
str(kid),
kpt,
colors=color,
font_sizes=self.radius * 3,
vertical_alignments='bottom',
horizontal_alignments='center')
return self.get_image()
def _draw_instance_heatmap(
self,
fields: PixelData,
overlaid_image: Optional[np.ndarray] = None,
):
"""Draw heatmaps of GT or prediction.
Args:
fields (:obj:`PixelData`): Data structure for
pixel-level annotations or predictions.
overlaid_image (np.ndarray): The image to draw.
Returns:
np.ndarray: the drawn image which channel is RGB.
"""
if 'heatmaps' not in fields:
return None
heatmaps = fields.heatmaps
if isinstance(heatmaps, np.ndarray):
heatmaps = torch.from_numpy(heatmaps)
if heatmaps.dim() == 3:
heatmaps, _ = heatmaps.max(dim=0)
heatmaps = heatmaps.unsqueeze(0)
out_image = self.draw_featmap(heatmaps, overlaid_image)
return out_image
def _draw_instance_xy_heatmap(
self,
fields: PixelData,
overlaid_image: Optional[np.ndarray] = None,
n: int = 20,
):
"""Draw heatmaps of GT or prediction.
Args:
fields (:obj:`PixelData`): Data structure for
pixel-level annotations or predictions.
overlaid_image (np.ndarray): The image to draw.
n (int): Number of keypoint, up to 20.
Returns:
np.ndarray: the drawn image which channel is RGB.
"""
if 'heatmaps' not in fields:
return None
heatmaps = fields.heatmaps
_, h, w = heatmaps.shape
if isinstance(heatmaps, np.ndarray):
heatmaps = torch.from_numpy(heatmaps)
out_image = SimCCVisualizer().draw_instance_xy_heatmap(
heatmaps, overlaid_image, n)
out_image = cv2.resize(out_image[:, :, ::-1], (w, h))
return out_image
@master_only
def add_datasample(self,
name: str,
image: np.ndarray,
data_sample: PoseDataSample,
draw_gt: bool = True,
draw_pred: bool = True,
draw_heatmap: bool = False,
draw_bbox: bool = False,
show_kpt_idx: bool = False,
skeleton_style: str = 'mmpose',
show: bool = False,
wait_time: float = 0,
out_file: Optional[str] = None,
kpt_thr: float = 0.3,
step: int = 0) -> None:
"""Draw datasample and save to all backends.
- If GT and prediction are plotted at the same time, they are
displayed in a stitched image where the left image is the
ground truth and the right image is the prediction.
- If ``show`` is True, all storage backends are ignored, and
the images will be displayed in a local window.
- If ``out_file`` is specified, the drawn image will be
saved to ``out_file``. t is usually used when the display
is not available.
Args:
name (str): The image identifier
image (np.ndarray): The image to draw
data_sample (:obj:`PoseDataSample`, optional): The data sample
to visualize
draw_gt (bool): Whether to draw GT PoseDataSample. Default to
``True``
draw_pred (bool): Whether to draw Prediction PoseDataSample.
Defaults to ``True``
draw_bbox (bool): Whether to draw bounding boxes. Default to
``False``
draw_heatmap (bool): Whether to draw heatmaps. Defaults to
``False``
show_kpt_idx (bool): Whether to show the index of keypoints.
Defaults to ``False``
skeleton_style (str): Skeleton style selection. Defaults to
``'mmpose'``
show (bool): Whether to display the drawn image. Default to
``False``
wait_time (float): The interval of show (s). Defaults to 0
out_file (str): Path to output file. Defaults to ``None``
kpt_thr (float, optional): Minimum threshold of keypoints
to be shown. Default: 0.3.
step (int): Global step value to record. Defaults to 0
"""
gt_img_data = None
pred_img_data = None
if draw_gt:
gt_img_data = image.copy()
gt_img_heatmap = None
# draw bboxes & keypoints
if 'gt_instances' in data_sample:
gt_img_data = self._draw_instances_kpts(
gt_img_data, data_sample.gt_instances, kpt_thr,
show_kpt_idx, skeleton_style)
if draw_bbox:
gt_img_data = self._draw_instances_bbox(
gt_img_data, data_sample.gt_instances)
# draw heatmaps
if 'gt_fields' in data_sample and draw_heatmap:
gt_img_heatmap = self._draw_instance_heatmap(
data_sample.gt_fields, image)
if gt_img_heatmap is not None:
gt_img_data = np.concatenate((gt_img_data, gt_img_heatmap),
axis=0)
if draw_pred:
pred_img_data = image.copy()
pred_img_heatmap = None
# draw bboxes & keypoints
if 'pred_instances' in data_sample:
pred_img_data = self._draw_instances_kpts(
pred_img_data, data_sample.pred_instances, kpt_thr,
show_kpt_idx, skeleton_style)
if draw_bbox:
pred_img_data = self._draw_instances_bbox(
pred_img_data, data_sample.pred_instances)
# draw heatmaps
if 'pred_fields' in data_sample and draw_heatmap:
if 'keypoint_x_labels' in data_sample.pred_instances:
pred_img_heatmap = self._draw_instance_xy_heatmap(
data_sample.pred_fields, image)
else:
pred_img_heatmap = self._draw_instance_heatmap(
data_sample.pred_fields, image)
if pred_img_heatmap is not None:
pred_img_data = np.concatenate(
(pred_img_data, pred_img_heatmap), axis=0)
# merge visualization results
if gt_img_data is not None and pred_img_data is not None:
if gt_img_heatmap is None and pred_img_heatmap is not None:
gt_img_data = np.concatenate((gt_img_data, image), axis=0)
elif gt_img_heatmap is not None and pred_img_heatmap is None:
pred_img_data = np.concatenate((pred_img_data, image), axis=0)
drawn_img = np.concatenate((gt_img_data, pred_img_data), axis=1)
elif gt_img_data is not None:
drawn_img = gt_img_data
else:
drawn_img = pred_img_data
# It is convenient for users to obtain the drawn image.
# For example, the user wants to obtain the drawn image and
# save it as a video during video inference.
self.set_image(drawn_img)
if show:
self.show(drawn_img, win_name=name, wait_time=wait_time)
if out_file is not None:
mmcv.imwrite(drawn_img[..., ::-1], out_file)
else:
# save drawn_img to backends
self.add_image(name, drawn_img, step)
return self.get_image()