Spaces:
Sleeping
Sleeping
| import cv2 | |
| from ultralytics.utils.plotting import colors | |
| class PlotPose: | |
| def __init__(self): | |
| self.skeleton = [ | |
| [16, 14], | |
| [14, 12], | |
| [17, 15], | |
| [15, 13], | |
| [12, 13], | |
| [6, 12], | |
| [7, 13], | |
| [6, 7], | |
| [6, 8], | |
| [7, 9], | |
| [8, 10], | |
| [9, 11], | |
| [2, 3], | |
| [1, 2], | |
| [1, 3], | |
| [2, 4], | |
| [3, 5], | |
| [4, 6], | |
| [5, 7], | |
| ] | |
| self.limb_color = colors.pose_palette[[9, 9, 9, 9, 7, 7, 7, 0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 16, 16]] | |
| self.kpt_color = colors.pose_palette[[16, 16, 16, 16, 16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9]] | |
| def plot(self,im, | |
| kpts, | |
| shape=(640, 640), | |
| radius=None, | |
| kpt_line=True, | |
| conf_thres= 0.25,): | |
| """ | |
| Plot keypoints on the image. | |
| Args: | |
| im (numpy array): img to plot keypoints | |
| kpts (numpy arrayr): Keypoints, shape [17, 3] (x, y, confidence). | |
| shape (tuple, optional): Image shape (h, w). | |
| radius (int, optional): Keypoint radius. | |
| kpt_line (bool, optional): Draw lines between keypoints. | |
| conf_thres (float, optional): Confidence threshold. | |
| kpt_color (tuple, optional): Keypoint color (B, G, R). | |
| Note: | |
| - `kpt_line=True` currently only supports human pose plotting. | |
| - Modifies self.im in-place. | |
| - If self.pil is True, converts image to numpy array and back to PIL. | |
| """ | |
| self.lw = 4 | |
| radius = radius if radius is not None else self.lw | |
| nkpt, ndim = kpts.shape | |
| is_pose = nkpt == 17 and ndim in {2, 3} | |
| kpt_line &= is_pose # `kpt_line=True` for now only supports human pose plotting | |
| for i, k in enumerate(kpts): | |
| color_k = self.kpt_color[i].tolist() | |
| # color_k = kpt_color or (kpt_color[i].tolist() if is_pose else colors(i)) | |
| x_coord, y_coord = k[0], k[1] | |
| if x_coord % shape[1] != 0 and y_coord % shape[0] != 0: | |
| if len(k) == 3: | |
| conf = k[2] | |
| if conf < conf_thres: | |
| continue | |
| cv2.circle(im, (int(x_coord), int(y_coord)), radius, 255, -1, lineType=cv2.LINE_AA) | |
| if kpt_line: | |
| ndim = kpts.shape[-1] | |
| for i, sk in enumerate(self.skeleton): | |
| pos1 = (int(kpts[(sk[0] - 1), 0]), int(kpts[(sk[0] - 1), 1])) | |
| pos2 = (int(kpts[(sk[1] - 1), 0]), int(kpts[(sk[1] - 1), 1])) | |
| if ndim == 3: | |
| conf1 = kpts[(sk[0] - 1), 2] | |
| conf2 = kpts[(sk[1] - 1), 2] | |
| if conf1 < conf_thres or conf2 < conf_thres: | |
| continue | |
| if pos1[0] % shape[1] == 0 or pos1[1] % shape[0] == 0 or pos1[0] < 0 or pos1[1] < 0: | |
| continue | |
| if pos2[0] % shape[1] == 0 or pos2[1] % shape[0] == 0 or pos2[0] < 0 or pos2[1] < 0: | |
| continue | |
| cv2.line( | |
| im, | |
| pos1, | |
| pos2, | |
| self.limb_color[i].tolist(), | |
| thickness=int(self.lw ), | |
| lineType=cv2.LINE_AA, | |
| ) | |
| return im | |