hmr2/utils/render_openpose.py

"""
Render OpenPose keypoints.
Code was ported to Python from the official C++ implementation https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/utilities/keypoint.cpp
"""
import cv2
import math
import numpy as np
from typing import List, Tuple

def get_keypoints_rectangle(keypoints: np.array, threshold: float) -> Tuple[float, float, float]:
 """
 Compute rectangle enclosing keypoints above the threshold.
 Args:
 keypoints (np.array): Keypoint array of shape (N, 3).
 threshold (float): Confidence visualization threshold.
 Returns:
 Tuple[float, float, float]: Rectangle width, height and area.
 """
 valid_ind = keypoints[:, -1] > threshold
 if valid_ind.sum() > 0:
 valid_keypoints = keypoints[valid_ind][:, :-1]
 max_x = valid_keypoints[:,0].max()
 max_y = valid_keypoints[:,1].max()
 min_x = valid_keypoints[:,0].min()
 min_y = valid_keypoints[:,1].min()
 width = max_x - min_x
 height = max_y - min_y
 area = width * height
 return width, height, area
 else:
 return 0,0,0

def render_keypoints(img: np.array,
 keypoints: np.array,
 pairs: List,
 colors: List,
 thickness_circle_ratio: float,
 thickness_line_ratio_wrt_circle: float,
 pose_scales: List,
 threshold: float = 0.1) -> np.array:
 """
 Render keypoints on input image.
 Args:
 img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
 keypoints (np.array): Keypoint array of shape (N, 3).
 pairs (List): List of keypoint pairs per limb.
 colors: (List): List of colors per keypoint.
 thickness_circle_ratio (float): Circle thickness ratio.
 thickness_line_ratio_wrt_circle (float): Line thickness ratio wrt the circle.
 pose_scales (List): List of pose scales.
 threshold (float): Only visualize keypoints with confidence above the threshold.
 Returns:
 (np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image. 
 """
 img_orig = img.copy()
 width, height = img.shape[1], img.shape[2]
 area = width * height

 lineType = 8
 shift = 0
 numberColors = len(colors)
 thresholdRectangle = 0.1

 person_width, person_height, person_area = get_keypoints_rectangle(keypoints, thresholdRectangle)
 if person_area > 0:
 ratioAreas = min(1, max(person_width / width, person_height / height))
 thicknessRatio = np.maximum(np.round(math.sqrt(area) * thickness_circle_ratio * ratioAreas), 2)
 thicknessCircle = np.maximum(1, thicknessRatio if ratioAreas > 0.05 else -np.ones_like(thicknessRatio))
 thicknessLine = np.maximum(1, np.round(thicknessRatio * thickness_line_ratio_wrt_circle))
 radius = thicknessRatio / 2

 img = np.ascontiguousarray(img.copy())
 for i, pair in enumerate(pairs):
 index1, index2 = pair
 if keypoints[index1, -1] > threshold and keypoints[index2, -1] > threshold:
 thicknessLineScaled = int(round(min(thicknessLine[index1], thicknessLine[index2]) * pose_scales[0]))
 colorIndex = index2
 color = colors[colorIndex % numberColors]
 keypoint1 = keypoints[index1, :-1].astype(np.int)
 keypoint2 = keypoints[index2, :-1].astype(np.int)
 cv2.line(img, tuple(keypoint1.tolist()), tuple(keypoint2.tolist()), tuple(color.tolist()), thicknessLineScaled, lineType, shift)
 for part in range(len(keypoints)):
 faceIndex = part
 if keypoints[faceIndex, -1] > threshold:
 radiusScaled = int(round(radius[faceIndex] * pose_scales[0]))
 thicknessCircleScaled = int(round(thicknessCircle[faceIndex] * pose_scales[0]))
 colorIndex = part
 color = colors[colorIndex % numberColors]
 center = keypoints[faceIndex, :-1].astype(np.int)
 cv2.circle(img, tuple(center.tolist()), radiusScaled, tuple(color.tolist()), thicknessCircleScaled, lineType, shift)
 return img

def render_body_keypoints(img: np.array,
 body_keypoints: np.array) -> np.array:
 """
 Render OpenPose body keypoints on input image.
 Args:
 img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
 body_keypoints (np.array): Keypoint array of shape (N, 3); 3 <====> (x, y, confidence).
 Returns:
 (np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image. 
 """

 thickness_circle_ratio = 1./75. * np.ones(body_keypoints.shape[0])
 thickness_line_ratio_wrt_circle = 0.75
 pairs = []
 pairs = [1,8,1,2,1,5,2,3,3,4,5,6,6,7,8,9,9,10,10,11,8,12,12,13,13,14,1,0,0,15,15,17,0,16,16,18,14,19,19,20,14,21,11,22,22,23,11,24]
 pairs = np.array(pairs).reshape(-1,2)
 colors = [255., 0., 85.,
 255., 0., 0.,
 255., 85., 0.,
 255., 170., 0.,
 255., 255., 0.,
 170., 255., 0.,
 85., 255., 0.,
 0., 255., 0.,
 255., 0., 0.,
 0., 255., 85.,
 0., 255., 170.,
 0., 255., 255.,
 0., 170., 255.,
 0., 85., 255.,
 0., 0., 255.,
 255., 0., 170.,
 170., 0., 255.,
 255., 0., 255.,
 85., 0., 255.,
 0., 0., 255.,
 0., 0., 255.,
 0., 0., 255.,
 0., 255., 255.,
 0., 255., 255.,
 0., 255., 255.]
 colors = np.array(colors).reshape(-1,3)
 pose_scales = [1]
 return render_keypoints(img, body_keypoints, pairs, colors, thickness_circle_ratio, thickness_line_ratio_wrt_circle, pose_scales, 0.1)

def render_openpose(img: np.array,
 body_keypoints: np.array) -> np.array:
 """
 Render keypoints in the OpenPose format on input image.
 Args:
 img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
 body_keypoints (np.array): Keypoint array of shape (N, 3); 3 <====> (x, y, confidence).
 Returns:
 (np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image. 
 """
 img = render_body_keypoints(img, body_keypoints)
 return img