|
""" |
|
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 |
|
|