mmpose-0.29.0/mmpose/core/visualization/image.py

# Copyright (c) OpenMMLab. All rights reserved.
import math
import os
import warnings

import cv2
import mmcv
import numpy as np
from matplotlib import pyplot as plt
from mmcv.utils.misc import deprecated_api_warning
from mmcv.visualization.color import color_val

try:
 import trimesh
 has_trimesh = True
except (ImportError, ModuleNotFoundError):
 has_trimesh = False

try:
 os.environ['PYOPENGL_PLATFORM'] = 'osmesa'
 import pyrender
 has_pyrender = True
except (ImportError, ModuleNotFoundError):
 has_pyrender = False


def imshow_bboxes(img,
 bboxes,
 labels=None,
 colors='green',
 text_color='white',
 thickness=1,
 font_scale=0.5,
 show=True,
 win_name='',
 wait_time=0,
 out_file=None):
 """Draw bboxes with labels (optional) on an image. This is a wrapper of
 mmcv.imshow_bboxes.

 Args:
 img (str or ndarray): The image to be displayed.
 bboxes (ndarray): ndarray of shape (k, 4), each row is a bbox in
 format [x1, y1, x2, y2].
 labels (str or list[str], optional): labels of each bbox.
 colors (list[str or tuple or :obj:`Color`]): A list of colors.
 text_color (str or tuple or :obj:`Color`): Color of texts.
 thickness (int): Thickness of lines.
 font_scale (float): Font scales of texts.
 show (bool): Whether to show the image.
 win_name (str): The window name.
 wait_time (int): Value of waitKey param.
 out_file (str, optional): The filename to write the image.

 Returns:
 ndarray: The image with bboxes drawn on it.
 """

 # adapt to mmcv.imshow_bboxes input format
 bboxes = np.split(
 bboxes, bboxes.shape[0], axis=0) if bboxes.shape[0] > 0 else []
 if not isinstance(colors, list):
 colors = [colors for _ in range(len(bboxes))]
 colors = [mmcv.color_val(c) for c in colors]
 assert len(bboxes) == len(colors)

 img = mmcv.imshow_bboxes(
 img,
 bboxes,
 colors,
 top_k=-1,
 thickness=thickness,
 show=False,
 out_file=None)

 if labels is not None:
 if not isinstance(labels, list):
 labels = [labels for _ in range(len(bboxes))]
 assert len(labels) == len(bboxes)

 for bbox, label, color in zip(bboxes, labels, colors):
 if label is None:
 continue
 bbox_int = bbox[0, :4].astype(np.int32)
 # roughly estimate the proper font size
 text_size, text_baseline = cv2.getTextSize(label,
 cv2.FONT_HERSHEY_DUPLEX,
 font_scale, thickness)
 text_x1 = bbox_int[0]
 text_y1 = max(0, bbox_int[1] - text_size[1] - text_baseline)
 text_x2 = bbox_int[0] + text_size[0]
 text_y2 = text_y1 + text_size[1] + text_baseline
 cv2.rectangle(img, (text_x1, text_y1), (text_x2, text_y2), color,
 cv2.FILLED)
 cv2.putText(img, label, (text_x1, text_y2 - text_baseline),
 cv2.FONT_HERSHEY_DUPLEX, font_scale,
 mmcv.color_val(text_color), thickness)

 if show:
 mmcv.imshow(img, win_name, wait_time)
 if out_file is not None:
 mmcv.imwrite(img, out_file)
 return img


@deprecated_api_warning({'pose_limb_color': 'pose_link_color'})
def imshow_keypoints(img,
 pose_result,
 skeleton=None,
 kpt_score_thr=0.3,
 pose_kpt_color=None,
 pose_link_color=None,
 radius=4,
 thickness=1,
 show_keypoint_weight=False):
 """Draw keypoints and links on an image.

 Args:
 img (str or Tensor): The image to draw poses on. If an image array
 is given, id will be modified in-place.
 pose_result (list[kpts]): The poses to draw. Each element kpts is
 a set of K keypoints as an Kx3 numpy.ndarray, where each
 keypoint is represented as x, y, score.
 kpt_score_thr (float, optional): Minimum score of keypoints
 to be shown. Default: 0.3.
 pose_kpt_color (np.array[Nx3]`): Color of N keypoints. If None,
 the keypoint will not be drawn.
 pose_link_color (np.array[Mx3]): Color of M links. If None, the
 links will not be drawn.
 thickness (int): Thickness of lines.
 """

 img = mmcv.imread(img)
 img_h, img_w, _ = img.shape

 for kpts in pose_result:

 kpts = np.array(kpts, copy=False)

 # draw each point on image
 if pose_kpt_color is not None:
 assert len(pose_kpt_color) == len(kpts)

 for kid, kpt in enumerate(kpts):
 x_coord, y_coord, kpt_score = int(kpt[0]), int(kpt[1]), kpt[2]

 if kpt_score < kpt_score_thr or pose_kpt_color[kid] is None:
 # skip the point that should not be drawn
 continue

 color = tuple(int(c) for c in pose_kpt_color[kid])
 if show_keypoint_weight:
 img_copy = img.copy()
 cv2.circle(img_copy, (int(x_coord), int(y_coord)), radius,
 color, -1)
 transparency = max(0, min(1, kpt_score))
 cv2.addWeighted(
 img_copy,
 transparency,
 img,
 1 - transparency,
 0,
 dst=img)
 else:
 cv2.circle(img, (int(x_coord), int(y_coord)), radius,
 color, -1)

 # draw links
 if skeleton is not None and pose_link_color is not None:
 assert len(pose_link_color) == len(skeleton)

 for sk_id, sk in enumerate(skeleton):
 pos1 = (int(kpts[sk[0], 0]), int(kpts[sk[0], 1]))
 pos2 = (int(kpts[sk[1], 0]), int(kpts[sk[1], 1]))

 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 kpts[sk[0], 2] < kpt_score_thr
 or kpts[sk[1], 2] < kpt_score_thr
 or pose_link_color[sk_id] is None):
 # skip the link that should not be drawn
 continue
 color = tuple(int(c) for c in pose_link_color[sk_id])
 if show_keypoint_weight:
 img_copy = img.copy()
 X = (pos1[0], pos2[0])
 Y = (pos1[1], pos2[1])
 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
 polygon = cv2.ellipse2Poly(
 (int(mX), int(mY)), (int(length / 2), int(stickwidth)),
 int(angle), 0, 360, 1)
 cv2.fillConvexPoly(img_copy, polygon, color)
 transparency = max(
 0, min(1, 0.5 * (kpts[sk[0], 2] + kpts[sk[1], 2])))
 cv2.addWeighted(
 img_copy,
 transparency,
 img,
 1 - transparency,
 0,
 dst=img)
 else:
 cv2.line(img, pos1, pos2, color, thickness=thickness)

 return img


def imshow_keypoints_3d(
 pose_result,
 img=None,
 skeleton=None,
 pose_kpt_color=None,
 pose_link_color=None,
 vis_height=400,
 kpt_score_thr=0.3,
 num_instances=-1,
 *,
 axis_azimuth=70,
 axis_limit=1.7,
 axis_dist=10.0,
 axis_elev=15.0,
):
 """Draw 3D keypoints and links in 3D coordinates.

 Args:
 pose_result (list[dict]): 3D pose results containing:
 - "keypoints_3d" ([K,4]): 3D keypoints
 - "title" (str): Optional. A string to specify the title of the
 visualization of this pose result
 img (str|np.ndarray): Opptional. The image or image path to show input
 image and/or 2D pose. Note that the image should be given in BGR
 channel order.
 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.ndarray[Nx3]`): Color of N keypoints. If None, do
 not nddraw keypoints.
 pose_link_color (np.array[Mx3]): Color of M links. If None, do not
 draw links.
 vis_height (int): The image height of the visualization. The width
 will be N*vis_height depending on the number of visualized
 items.
 kpt_score_thr (float): Minimum score of keypoints to be shown.
 Default: 0.3.
 num_instances (int): Number of instances to be shown in 3D. If smaller
 than 0, all the instances in the pose_result will be shown.
 Otherwise, pad or truncate the pose_result to a length of
 num_instances.
 axis_azimuth (float): axis azimuth angle for 3D visualizations.
 axis_dist (float): axis distance for 3D visualizations.
 axis_elev (float): axis elevation view angle for 3D visualizations.
 axis_limit (float): The axis limit to visualize 3d pose. The xyz
 range will be set as:
 - x: [x_c - axis_limit/2, x_c + axis_limit/2]
 - y: [y_c - axis_limit/2, y_c + axis_limit/2]
 - z: [0, axis_limit]
 Where x_c, y_c is the mean value of x and y coordinates
 figsize: (float): figure size in inch.
 """

 show_img = img is not None
 if num_instances < 0:
 num_instances = len(pose_result)
 else:
 if len(pose_result) > num_instances:
 pose_result = pose_result[:num_instances]
 elif len(pose_result) < num_instances:
 pose_result += [dict()] * (num_instances - len(pose_result))
 num_axis = num_instances + 1 if show_img else num_instances

 plt.ioff()
 fig = plt.figure(figsize=(vis_height * num_axis * 0.01, vis_height * 0.01))

 if show_img:
 img = mmcv.imread(img, channel_order='bgr')
 img = mmcv.bgr2rgb(img)
 img = mmcv.imrescale(img, scale=vis_height / img.shape[0])

 ax_img = fig.add_subplot(1, num_axis, 1)
 ax_img.get_xaxis().set_visible(False)
 ax_img.get_yaxis().set_visible(False)
 ax_img.set_axis_off()
 ax_img.set_title('Input')
 ax_img.imshow(img, aspect='equal')

 for idx, res in enumerate(pose_result):
 dummy = len(res) == 0
 kpts = np.zeros((1, 3)) if dummy else res['keypoints_3d']
 if kpts.shape[1] == 3:
 kpts = np.concatenate([kpts, np.ones((kpts.shape[0], 1))], axis=1)
 valid = kpts[:, 3] >= kpt_score_thr

 ax_idx = idx + 2 if show_img else idx + 1
 ax = fig.add_subplot(1, num_axis, ax_idx, projection='3d')
 ax.view_init(
 elev=axis_elev,
 azim=axis_azimuth,
 )
 x_c = np.mean(kpts[valid, 0]) if sum(valid) > 0 else 0
 y_c = np.mean(kpts[valid, 1]) if sum(valid) > 0 else 0
 ax.set_xlim3d([x_c - axis_limit / 2, x_c + axis_limit / 2])
 ax.set_ylim3d([y_c - axis_limit / 2, y_c + axis_limit / 2])
 ax.set_zlim3d([0, axis_limit])
 ax.set_aspect('auto')
 ax.set_xticks([])
 ax.set_yticks([])
 ax.set_zticks([])
 ax.set_xticklabels([])
 ax.set_yticklabels([])
 ax.set_zticklabels([])
 ax.dist = axis_dist

 if not dummy and pose_kpt_color is not None:
 pose_kpt_color = np.array(pose_kpt_color)
 assert len(pose_kpt_color) == len(kpts)
 x_3d, y_3d, z_3d = np.split(kpts[:, :3], [1, 2], axis=1)
 # matplotlib uses RGB color in [0, 1] value range
 _color = pose_kpt_color[..., ::-1] / 255.
 ax.scatter(
 x_3d[valid],
 y_3d[valid],
 z_3d[valid],
 marker='o',
 color=_color[valid],
 )

 if not dummy and skeleton is not None and pose_link_color is not None:
 pose_link_color = np.array(pose_link_color)
 assert len(pose_link_color) == len(skeleton)
 for link, link_color in zip(skeleton, pose_link_color):
 link_indices = [_i for _i in link]
 xs_3d = kpts[link_indices, 0]
 ys_3d = kpts[link_indices, 1]
 zs_3d = kpts[link_indices, 2]
 kpt_score = kpts[link_indices, 3]
 if kpt_score.min() > kpt_score_thr:
 # matplotlib uses RGB color in [0, 1] value range
 _color = link_color[::-1] / 255.
 ax.plot(xs_3d, ys_3d, zs_3d, color=_color, zdir='z')

 if 'title' in res:
 ax.set_title(res['title'])

 # convert figure to numpy array
 fig.tight_layout()
 fig.canvas.draw()
 img_w, img_h = fig.canvas.get_width_height()
 img_vis = np.frombuffer(
 fig.canvas.tostring_rgb(), dtype=np.uint8).reshape(img_h, img_w, -1)
 img_vis = mmcv.rgb2bgr(img_vis)

 plt.close(fig)

 return img_vis


def imshow_mesh_3d(img,
 vertices,
 faces,
 camera_center,
 focal_length,
 colors=(76, 76, 204)):
 """Render 3D meshes on background image.

 Args:
 img(np.ndarray): Background image.
 vertices (list of np.ndarray): Vetrex coordinates in camera space.
 faces (list of np.ndarray): Faces of meshes.
 camera_center ([2]): Center pixel.
 focal_length ([2]): Focal length of camera.
 colors (list[str or tuple or Color]): A list of mesh colors.
 """

 H, W, C = img.shape

 if not has_pyrender:
 warnings.warn('pyrender package is not installed.')
 return img

 if not has_trimesh:
 warnings.warn('trimesh package is not installed.')
 return img

 try:
 renderer = pyrender.OffscreenRenderer(
 viewport_width=W, viewport_height=H)
 except (ImportError, RuntimeError):
 warnings.warn('pyrender package is not installed correctly.')
 return img

 if not isinstance(colors, list):
 colors = [colors for _ in range(len(vertices))]
 colors = [color_val(c) for c in colors]

 depth_map = np.ones([H, W]) * np.inf
 output_img = img
 for idx in range(len(vertices)):
 color = colors[idx]
 color = [c / 255.0 for c in color]
 color.append(1.0)
 vert = vertices[idx]
 face = faces[idx]

 material = pyrender.MetallicRoughnessMaterial(
 metallicFactor=0.2, alphaMode='OPAQUE', baseColorFactor=color)

 mesh = trimesh.Trimesh(vert, face)
 rot = trimesh.transformations.rotation_matrix(
 np.radians(180), [1, 0, 0])
 mesh.apply_transform(rot)
 mesh = pyrender.Mesh.from_trimesh(mesh, material=material)

 scene = pyrender.Scene(ambient_light=(0.5, 0.5, 0.5))
 scene.add(mesh, 'mesh')

 camera_pose = np.eye(4)
 camera = pyrender.IntrinsicsCamera(
 fx=focal_length[0],
 fy=focal_length[1],
 cx=camera_center[0],
 cy=camera_center[1],
 zfar=1e5)
 scene.add(camera, pose=camera_pose)

 light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
 light_pose = np.eye(4)

 light_pose[:3, 3] = np.array([0, -1, 1])
 scene.add(light, pose=light_pose)

 light_pose[:3, 3] = np.array([0, 1, 1])
 scene.add(light, pose=light_pose)

 light_pose[:3, 3] = np.array([1, 1, 2])
 scene.add(light, pose=light_pose)

 color, rend_depth = renderer.render(
 scene, flags=pyrender.RenderFlags.RGBA)

 valid_mask = (rend_depth < depth_map) * (rend_depth > 0)
 depth_map[valid_mask] = rend_depth[valid_mask]
 valid_mask = valid_mask[:, :, None]
 output_img = (
 valid_mask * color[:, :, :3] + (1 - valid_mask) * output_img)

 return output_img


def imshow_multiview_keypoints_3d(
 pose_result,
 skeleton=None,
 pose_kpt_color=None,
 pose_link_color=None,
 space_size=[8000, 8000, 2000],
 space_center=[0, -500, 800],
 kpt_score_thr=0.0,
):
 """Draw 3D keypoints and links in 3D coordinates.

 Args:
 pose_result (list[kpts]): The poses to draw. Each element kpts is
 a set of K keypoints as an Kx4 numpy.ndarray, where each
 keypoint is represented as x, y, z, score.
 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.ndarray[Nx3]`): Color of N keypoints. If None, do
 not nddraw keypoints.
 pose_link_color (np.array[Mx3]): Color of M links. If None, do not
 draw links.
 space_size: (list). Default: [8000, 8000, 2000].
 space_center: (list). Default: [0, -500, 800].
 kpt_score_thr (float): Minimum score of keypoints to be shown.
 Default: 0.0.
 """
 fig = plt.figure()
 ax = plt.axes(projection='3d')
 ax.set_xlim3d(space_center[0] - space_size[0] * 0.5,
 space_center[0] + space_size[0] * 0.5)
 ax.set_ylim3d(space_center[1] - space_size[1] * 0.5,
 space_center[1] + space_size[1] * 0.5)
 ax.set_zlim3d(space_center[2] - space_size[2] * 0.5,
 space_center[2] + space_size[2] * 0.5)
 pose_kpt_color = np.array(pose_kpt_color)
 pose_kpt_color = pose_kpt_color[..., ::-1] / 255.

 for kpts in pose_result:
 # draw each point on image
 xs, ys, zs, scores = kpts.T
 valid = scores > kpt_score_thr
 ax.scatter(
 xs[valid],
 ys[valid],
 zs[valid],
 marker='o',
 color=pose_kpt_color[valid])

 for link, link_color in zip(skeleton, pose_link_color):
 link_indices = [_i for _i in link]
 xs_3d = kpts[link_indices, 0]
 ys_3d = kpts[link_indices, 1]
 zs_3d = kpts[link_indices, 2]
 kpt_score = kpts[link_indices, 3]
 if kpt_score.min() > kpt_score_thr:
 # matplotlib uses RGB color in [0, 1] value range
 _color = np.array(link_color[::-1]) / 255.
 ax.plot(xs_3d, ys_3d, zs_3d, color=_color)

 # convert figure to numpy array
 fig.tight_layout()
 fig.canvas.draw()
 img_w, img_h = fig.canvas.get_width_height()
 img_vis = np.frombuffer(
 fig.canvas.tostring_rgb(), dtype=np.uint8).reshape(img_h, img_w, -1)
 img_vis = mmcv.rgb2bgr(img_vis)

 plt.close(fig)

 return img_vis