RTMPose without mmcv, mmpose, mmdet etc.
from .draw import draw_bbox, draw_skeleton
from .tools import YOLOX, Body, PoseTracker, RTMDet, RTMPose, Wholebody

__all__ = [
    'RTMDet', 'RTMPose', 'YOLOX', 'Wholebody', 'Body', 'draw_skeleton',
    'draw_bbox', 'PoseTracker'
]

rtmlib/__init__.py
Tau-J/rtmlib
import math

import cv2
import numpy as np

from .visualization.skeleton import *  # noqa


def draw_bbox(img, bboxes, color=(0, 255, 0)):
    for bbox in bboxes:
        img = cv2.rectangle(img, (int(bbox[0]), int(bbox[1])),
                            (int(bbox[2]), int(bbox[3])), color, 2)
    return img


def draw_skeleton(img,
                  keypoints,
                  scores,
                  openpose_skeleton=False,
                  kpt_thr=0.5,
                  radius=2,
                  line_width=2):
    if openpose_skeleton:
        skeleton = 'openpose18' if keypoints.shape[1] == 18 else 'openpose134'
    else:
        skeleton = 'coco17' if keypoints.shape[1] == 17 else 'coco133'

    skeleton_dict = eval(f'{skeleton}')
    keypoint_info = skeleton_dict['keypoint_info']
    skeleton_info = skeleton_dict['skeleton_info']

    if len(keypoints.shape) == 2:
        keypoints = keypoints[None, :, :]
        scores = scores[None, :, :]

    num_instance = keypoints.shape[0]
    if skeleton.startswith('coco'):
        for i in range(num_instance):
            img = draw_mmpose(img, keypoints[i], scores[i], keypoint_info,
                              skeleton_info, kpt_thr, radius, line_width)
    else:
        for i in range(num_instance):
            img = draw_openpose(img,
                                keypoints[i],
                                scores[i],
                                keypoint_info,
                                skeleton_info,
                                kpt_thr,
                                radius * 2,
                                alpha=0.6,
                                line_width=line_width * 2)
    return img


def draw_mmpose(img,
                keypoints,
                scores,
                keypoint_info,
                skeleton_info,
                kpt_thr=0.5,
                radius=2,
                line_width=2):
    assert len(keypoints.shape) == 2

    vis_kpt = [s >= kpt_thr for s in scores]

    link_dict = {}
    for i, kpt_info in keypoint_info.items():
        kpt_color = tuple(kpt_info['color'])
        link_dict[kpt_info['name']] = kpt_info['id']

        kpt = keypoints[i]

        if vis_kpt[i]:
            img = cv2.circle(img, (int(kpt[0]), int(kpt[1])), int(radius),
                             kpt_color, -1)

    for i, ske_info in skeleton_info.items():
        link = ske_info['link']
        pt0, pt1 = link_dict[link[0]], link_dict[link[1]]

        if vis_kpt[pt0] and vis_kpt[pt1]:
            link_color = ske_info['color']
            kpt0 = keypoints[pt0]
            kpt1 = keypoints[pt1]

            img = cv2.line(img, (int(kpt0[0]), int(kpt0[1])),
                           (int(kpt1[0]), int(kpt1[1])),
                           link_color,
                           thickness=line_width)

    return img


def draw_openpose(img,
                  keypoints,
                  scores,
                  keypoint_info,
                  skeleton_info,
                  kpt_thr=0.4,
                  radius=4,
                  alpha=1.0,
                  line_width=2):
    h, w = img.shape[:2]

    link_dict = {}
    for i, kpt_info in keypoint_info.items():
        kpt_color = tuple(kpt_info['color'])
        link_dict[kpt_info['name']] = kpt_info['id']

    for i, ske_info in skeleton_info.items():
        link = ske_info['link']
        pt0, pt1 = link_dict[link[0]], link_dict[link[1]]

        link_color = ske_info['color']
        kpt0, kpt1 = keypoints[pt0], keypoints[pt1]
        s0, s1 = scores[pt0], scores[pt1]

        if (kpt0[0] <= 0 or kpt0[1] >= w or kpt0[1] <= 0 or kpt0[1] >= h
                or kpt1[0] <= 0 or kpt1[1] >= w or kpt1[1] <= 0 or kpt1[1] >= h
                or s0 < kpt_thr or s1 < kpt_thr or link_color is None):
            continue

        X = np.array([kpt0[0], kpt1[0]])
        Y = np.array([kpt0[1], kpt1[1]])

        if i <= 16:
            # body part
            mX = np.mean(X)
            mY = np.mean(Y)
            length = ((Y[0] - Y[1])**2 + (X[0] - X[1])**2)**0.5
            transparency = 0.6
            angle = math.degrees(math.atan2(Y[0] - Y[1], X[0] - X[1]))
            polygons = cv2.ellipse2Poly((int(mX), int(mY)),
                                        (int(length / 2), int(line_width)),
                                        int(angle), 0, 360, 1)
            img = draw_polygons(img,
                                polygons,
                                edge_colors=link_color,
                                alpha=transparency)
        else:
            img = cv2.line(img, (int(X[0]), int(Y[0])), (int(X[1]), int(Y[1])),
                           link_color,
                           thickness=2)

    for j, kpt_info in keypoint_info.items():
        kpt_color = tuple(kpt_info['color'][::-1])
        kpt = keypoints[j]

        if scores[j] < kpt_thr or sum(kpt_color) == 0:
            continue

        transparency = alpha
        if 24 <= j <= 91:
            j_radius = 3
        else:
            j_radius = 4
        # j_radius = radius // 2 if j > 17 else radius

        img = draw_circles(img,
                           kpt,
                           radius=np.array([j_radius]),
                           face_colors=kpt_color,
                           alpha=transparency)

    return img


def draw_polygons(img, polygons, edge_colors, alpha=1.0):
    if alpha == 1.0:
        img = cv2.fillConvexPoly(img, polygons, edge_colors)
    else:
        img = cv2.fillConvexPoly(img.copy(), polygons, edge_colors)
        img = cv2.addWeighted(img, 1 - alpha, img, alpha, 0)
    return img


def draw_circles(img, center, radius, face_colors, alpha=1.0):
    if alpha == 1.0:
        img = cv2.circle(img, (int(center[0]), int(center[1])), int(radius),
                         face_colors, -1)
    else:
        img = cv2.circle(img.copy(), (int(center[0]), int(center[1])),
                         int(radius), face_colors, -1)
        img = cv2.addWeighted(img, 1 - alpha, img, alpha, 0)
    return img

rtmlib/draw.py
Tau-J/rtmlib
from .object_detection import YOLOX, RTMDet
from .pose_estimation import RTMPose
from .solution import Body, PoseTracker, Wholebody

__all__ = ['RTMDet', 'RTMPose', 'YOLOX', 'Wholebody', 'Body', 'PoseTracker']

rtmlib/tools/__init__.py
Tau-J/rtmlib
import os
from abc import ABCMeta, abstractmethod
from typing import Any

import cv2
import numpy as np

from .file import download_checkpoint

RTMLIB_SETTINGS = {
    'opencv': {
        'cpu': (cv2.dnn.DNN_BACKEND_OPENCV, cv2.dnn.DNN_TARGET_CPU),

        # You need to manually build OpenCV through cmake
        'cuda': (cv2.dnn.DNN_BACKEND_CUDA, cv2.dnn.DNN_TARGET_CUDA)
    },
    'onnxruntime': {
        'cpu': 'CPUExecutionProvider',
        'cuda': 'CUDAExecutionProvider'
    },
}


class BaseTool(metaclass=ABCMeta):

    def __init__(self,
                 onnx_model: str = None,
                 model_input_size: tuple = None,
                 mean: tuple = None,
                 std: tuple = None,
                 backend: str = 'opencv',
                 device: str = 'cpu'):

        if not os.path.exists(onnx_model):
            onnx_model = download_checkpoint(onnx_model)

        if backend == 'opencv':
            try:
                providers = RTMLIB_SETTINGS[backend][device]

                session = cv2.dnn.readNetFromONNX(onnx_model)
                session.setPreferableBackend(providers[0])
                session.setPreferableTarget(providers[1])
                self.session = session
            except Exception:
                raise RuntimeError(
                    'This model is not supported by OpenCV'
                    ' backend, please use `pip install'
                    ' onnxruntime` or `pip install'
                    ' onnxruntime-gpu` to install onnxruntime'
                    ' backend. Then specify `backend=onnxruntime`.')  # noqa

        elif backend == 'onnxruntime':
            import onnxruntime as ort
            providers = RTMLIB_SETTINGS[backend][device]

            self.session = ort.InferenceSession(path_or_bytes=onnx_model,
                                                providers=[providers])

        elif backend == 'openvino':
            from openvino.runtime import Core
            core = Core()
            model_onnx = core.read_model(model=onnx_model)

            if device != 'cpu':
                print('OpenVINO only supports CPU backend, automatically'
                      ' switched to CPU backend.')

            self.compiled_model = core.compile_model(
                model=model_onnx,
                device_name='CPU',
                config={'PERFORMANCE_HINT': 'LATENCY'})
            self.input_layer = self.compiled_model.input(0)
            self.output_layer0 = self.compiled_model.output(0)
            self.output_layer1 = self.compiled_model.output(1)

        else:
            raise NotImplementedError

        print(f'load {onnx_model} with {backend} backend')

        self.onnx_model = onnx_model
        self.model_input_size = model_input_size
        self.mean = mean
        self.std = std
        self.backend = backend
        self.device = device

    @abstractmethod
    def __call__(self, *args, **kwargs) -> Any:
        """Implement the actual function here."""
        raise NotImplementedError

    def inference(self, img: np.ndarray):
        """Inference model.

        Args:
            img (np.ndarray): Input image in shape.

        Returns:
            outputs (np.ndarray): Output of RTMPose model.
        """
        # build input to (1, 3, H, W)
        img = img.transpose(2, 0, 1)
        img = np.ascontiguousarray(img, dtype=np.float32)
        input = img[None, :, :, :]

        # run model
        if self.backend == 'opencv':
            outNames = self.session.getUnconnectedOutLayersNames()
            self.session.setInput(input)
            outputs = self.session.forward(outNames)
        elif self.backend == 'onnxruntime':
            sess_input = {self.session.get_inputs()[0].name: input}
            sess_output = []
            for out in self.session.get_outputs():
                sess_output.append(out.name)

            outputs = self.session.run(sess_output, sess_input)
        elif self.backend == 'openvino':
            results = self.compiled_model(input)
            output0 = results[self.output_layer0]
            output1 = results[self.output_layer1]
            if output0.ndim == output1.ndim:
                outputs = [output0, output1]
            else:
                outputs = [output0]

        return outputs

rtmlib/tools/base.py
Tau-J/rtmlib
import hashlib
import os
import re
import sys
import tempfile
from glob import glob
from pathlib import Path
from typing import Optional
from urllib.parse import urlparse  # noqa: F401
from urllib.request import Request, urlopen

from tqdm import tqdm


def _get_rtmhub_dir():
    torch_home = os.path.expanduser(
        os.getenv(
            'TORCH_HOME',
            os.path.join(os.getenv('XDG_CACHE_HOME', '~/.cache'), 'rtmlib')))
    return os.path.join(torch_home, 'hub')


def download_url_to_file(url, dst, hash_prefix=None, progress=True):
    """Download object at the given URL to a local path.

    Modified from `torch.hub.download_url_to_file`.

    Args:
        url (str): URL of the object to download
        dst (str): Full path where object will be saved, for example,
            ``/tmp/temporary_file``.
        hash_prefix (str, optional): If not None, the SHA256 downloaded
            file should start with ``hash_prefix``. Defaults to None.
        progress (bool): whether or not to display a progress
            bar to stderr Defaults to True.
    """
    file_size = None
    req = Request(url, headers={'User-Agent': 'mmlmtools'})
    u = urlopen(req)
    meta = u.info()
    if hasattr(meta, 'getheaders'):
        content_length = meta.getheaders('Content-Length')
    else:
        content_length = meta.get_all('Content-Length')
    if content_length is not None and len(content_length) > 0:
        file_size = int(content_length[0])

    # We deliberately save it in a temp file and move it after
    # download is complete. This prevents a local working checkpoint
    # being overridden by a broken download.
    dst = Path(dst).expanduser().absolute()
    f = tempfile.NamedTemporaryFile(delete=False, dir=dst.parent)

    try:
        if hash_prefix is not None:
            sha256 = hashlib.sha256()
        with tqdm(total=file_size,
                  disable=not progress,
                  unit='B',
                  unit_scale=True,
                  unit_divisor=1024) as pbar:
            while True:
                buffer = u.read(8192)
                if len(buffer) == 0:
                    break
                f.write(buffer)
                if hash_prefix is not None:
                    sha256.update(buffer)
                pbar.update(len(buffer))

        f.close()
        if hash_prefix is not None:
            digest = sha256.hexdigest()
            if digest[:len(hash_prefix)] != hash_prefix:
                raise RuntimeError(
                    'invalid hash value (expected "{}", got "{}")'.format(
                        hash_prefix, digest))
        Path(f.name).rename(dst)
    finally:
        f.close()
        if os.path.exists(f.name):
            os.remove(f.name)


def download_checkpoint(url: str,
                        dst_dir: Optional[str] = None,
                        progress: bool = True,
                        check_hash: bool = False,
                        filename: Optional[str] = None) -> str:
    """Download the checkpoint from the given URL.

    Modified from `torch.hub.load_state_dict_from_url`.

    If the object is already present in `dst_dir`, it will be returned
    directly.
    The default value of ``dst_dir`` is the same as the checkpoint cache
    path of PyTorch hub.

    Args:
        url (str): URL of the object to download
        model_dir (str, optional): directory in which to save the object
        progress (bool, optional): whether or not to display a progress bar to
            stderr. Defaults to True.
        check_hash(bool, optional): If True, the filename part of the URL
            should follow the naming convention ``filename-<sha256>.ext`` where
            ``<sha256>`` is the first eight or more digits of the SHA256 hash
            of the contents of the file. The hash is used to ensure unique
            names and to verify the contents of the file. Defaults to False.
        filename (str, optional): name for the downloaded file.
            Filename from ``url`` will be used if not set.

    Returns:
        str: The path of the downloaded file.
    """
    if dst_dir is None:
        dst_dir = os.path.join(_get_rtmhub_dir(), 'checkpoints')

    dst_dir = Path(dst_dir)
    dst_dir.mkdir(parents=True, exist_ok=True)

    parts = urlparse(url)
    filename = filename or os.path.basename(parts.path)
    cached_file = dst_dir / filename
    onnx_name = f"{dst_dir}/{str(filename).split('.')[0]}.onnx"

    if not cached_file.exists():
        if os.path.exists(onnx_name):
            return onnx_name

        sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
        hash_prefix = None
        if check_hash:
            HASH_REGEX = re.compile(r'-([a-f0-9]*)\.')
            r = HASH_REGEX.search(filename)  # r is Optional[Match[str]]
            hash_prefix = r.group(1) if r else None
        download_url_to_file(url, cached_file, hash_prefix, progress=progress)

    if str(cached_file).split('.')[-1] == 'zip':
        os.system(f'unzip -d {dst_dir}/tmp {cached_file}')
        cached_list = glob(f'{dst_dir}/**', recursive=True)

        for each in cached_list:
            if each[-12:] == 'end2end.onnx':
                cached_onnx = each
                break
        os.system(f'mv {cached_onnx} {onnx_name}')
        os.system(f'rm -rf {cached_file}')
        os.system(f'rm -rf {dst_dir}/tmp')
        cached_file = onnx_name

    return str(cached_file)

rtmlib/tools/file.py
Tau-J/rtmlib
from .rtmdet import RTMDet
from .yolox import YOLOX

__all__ = ['RTMDet', 'YOLOX']

rtmlib/tools/object_detection/__init__.py
Tau-J/rtmlib
import numpy as np


def nms(boxes, scores, nms_thr):
    """Single class NMS implemented in Numpy."""
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= nms_thr)[0]
        order = order[inds + 1]

    return keep


def multiclass_nms(boxes, scores, nms_thr, score_thr):
    """Multiclass NMS implemented in Numpy.

    Class-aware version.
    """
    final_dets = []
    num_classes = scores.shape[1]
    for cls_ind in range(num_classes):
        cls_scores = scores[:, cls_ind]
        valid_score_mask = cls_scores > score_thr
        if valid_score_mask.sum() == 0:
            continue
        else:
            valid_scores = cls_scores[valid_score_mask]
            valid_boxes = boxes[valid_score_mask]
            keep = nms(valid_boxes, valid_scores, nms_thr)
            if len(keep) > 0:
                cls_inds = np.ones((len(keep), 1)) * cls_ind
                dets = np.concatenate(
                    [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1)
                final_dets.append(dets)
    if len(final_dets) == 0:
        return None
    return np.concatenate(final_dets, 0)

rtmlib/tools/object_detection/post_processings.py
Tau-J/rtmlib
from typing import List, Tuple

import cv2
import numpy as np

from ..base import BaseTool
from .post_processings import multiclass_nms


class RTMDet(BaseTool):

    def __init__(self,
                 onnx_model: str,
                 model_input_size: tuple = (640, 640),
                 backend: str = 'onnxruntime',
                 device: str = 'cpu'):
        super().__init__(onnx_model, model_input_size, backend, device)

    def __call__(self, image: np.ndarray):
        image, ratio = self.preprocess(image)
        outputs = self.inference(image)
        results = self.postprocess(outputs, ratio)
        return results

    def preprocess(self, img: np.ndarray):
        """Do preprocessing for RTMPose model inference.

        Args:
            img (np.ndarray): Input image in shape.

        Returns:
            tuple:
            - resized_img (np.ndarray): Preprocessed image.
            - center (np.ndarray): Center of image.
            - scale (np.ndarray): Scale of image.
        """
        if len(img.shape) == 3:
            padded_img = np.ones(
                (self.model_input_size[0], self.model_input_size[1], 3),
                dtype=np.uint8) * 114
        else:
            padded_img = np.ones(self.model_input_size, dtype=np.uint8) * 114

        ratio = min(self.model_input_size[0] / img.shape[0],
                    self.model_input_size[1] / img.shape[1])
        resized_img = cv2.resize(
            img,
            (int(img.shape[1] * ratio), int(img.shape[0] * ratio)),
            interpolation=cv2.INTER_LINEAR,
        ).astype(np.uint8)
        padded_shape = (int(img.shape[0] * ratio), int(img.shape[1] * ratio))
        padded_img[:padded_shape[0], :padded_shape[1]] = resized_img

        return padded_img, ratio

    def postprocess(
        self,
        outputs: List[np.ndarray],
        ratio: float = 1.,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Do postprocessing for RTMPose model inference.

        Args:
            outputs (List[np.ndarray]): Outputs of RTMPose model.
            ratio (float): Ratio of preprocessing.

        Returns:
            tuple:
            - final_boxes (np.ndarray): Final bounding boxes.
            - final_scores (np.ndarray): Final scores.
        """

        if outputs.shape[-1] == 4:
            # onnx without nms module

            grids = []
            expanded_strides = []
            strides = [8, 16, 32]

            hsizes = [self.model_input_size[0] // stride for stride in strides]
            wsizes = [self.model_input_size[1] // stride for stride in strides]

            for hsize, wsize, stride in zip(hsizes, wsizes, strides):
                xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
                grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
                grids.append(grid)
                shape = grid.shape[:2]
                expanded_strides.append(np.full((*shape, 1), stride))

            grids = np.concatenate(grids, 1)
            expanded_strides = np.concatenate(expanded_strides, 1)
            outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
            outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides

            predictions = outputs[0]
            boxes = predictions[:, :4]
            scores = predictions[:, 4:5] * predictions[:, 5:]

            boxes_xyxy = np.ones_like(boxes)
            boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.
            boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.
            boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.
            boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.
            boxes_xyxy /= ratio
            dets = multiclass_nms(boxes_xyxy,
                                  scores,
                                  nms_thr=self.nms_thr,
                                  score_thr=self.score_thr)
            if dets is not None:
                pack_dets = (dets[:, :4], dets[:, 4], dets[:, 5])
                final_boxes, final_scores, final_cls_inds = pack_dets
                isscore = final_scores > 0.3
                iscat = final_cls_inds == 0
                isbbox = [i and j for (i, j) in zip(isscore, iscat)]
                final_boxes = final_boxes[isbbox]

        elif outputs.shape[-1] == 5:
            # onnx contains nms module

            pack_dets = (outputs[0, :, :4], outputs[0, :, 4])
            final_boxes, final_scores = pack_dets
            final_boxes /= ratio
            isscore = final_scores > 0.3
            isbbox = [i for i in isscore]
            final_boxes = final_boxes[isbbox]

        return final_boxes

rtmlib/tools/object_detection/rtmdet.py
Tau-J/rtmlib
# Code modified from https://github.com/IDEA-Research/DWPose/blob/opencv_onnx/ControlNet-v1-1-nightly/annotator/dwpose/cv_ox_det.py  # noqa
from typing import List, Tuple

import cv2
import numpy as np

from ..base import BaseTool
from .post_processings import multiclass_nms


class YOLOX(BaseTool):

    def __init__(self,
                 onnx_model: str,
                 model_input_size: tuple = (640, 640),
                 nms_thr=0.45,
                 score_thr=0.7,
                 backend: str = 'onnxruntime',
                 device: str = 'cpu'):
        super().__init__(onnx_model,
                         model_input_size,
                         backend=backend,
                         device=device)
        self.nms_thr = nms_thr
        self.score_thr = score_thr

    def __call__(self, image: np.ndarray):
        image, ratio = self.preprocess(image)
        outputs = self.inference(image)[0]
        results = self.postprocess(outputs, ratio)
        return results

    def preprocess(self, img: np.ndarray):
        """Do preprocessing for RTMPose model inference.

        Args:
            img (np.ndarray): Input image in shape.

        Returns:
            tuple:
            - resized_img (np.ndarray): Preprocessed image.
            - center (np.ndarray): Center of image.
            - scale (np.ndarray): Scale of image.
        """
        if len(img.shape) == 3:
            padded_img = np.ones(
                (self.model_input_size[0], self.model_input_size[1], 3),
                dtype=np.uint8) * 114
        else:
            padded_img = np.ones(self.model_input_size, dtype=np.uint8) * 114

        ratio = min(self.model_input_size[0] / img.shape[0],
                    self.model_input_size[1] / img.shape[1])
        resized_img = cv2.resize(
            img,
            (int(img.shape[1] * ratio), int(img.shape[0] * ratio)),
            interpolation=cv2.INTER_LINEAR,
        ).astype(np.uint8)
        padded_shape = (int(img.shape[0] * ratio), int(img.shape[1] * ratio))
        padded_img[:padded_shape[0], :padded_shape[1]] = resized_img

        return padded_img, ratio

    def postprocess(
        self,
        outputs: List[np.ndarray],
        ratio: float = 1.,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Do postprocessing for RTMPose model inference.

        Args:
            outputs (List[np.ndarray]): Outputs of RTMPose model.
            ratio (float): Ratio of preprocessing.

        Returns:
            tuple:
            - final_boxes (np.ndarray): Final bounding boxes.
            - final_scores (np.ndarray): Final scores.
        """

        if outputs.shape[-1] == 4:
            # onnx without nms module

            grids = []
            expanded_strides = []
            strides = [8, 16, 32]

            hsizes = [self.model_input_size[0] // stride for stride in strides]
            wsizes = [self.model_input_size[1] // stride for stride in strides]

            for hsize, wsize, stride in zip(hsizes, wsizes, strides):
                xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
                grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
                grids.append(grid)
                shape = grid.shape[:2]
                expanded_strides.append(np.full((*shape, 1), stride))

            grids = np.concatenate(grids, 1)
            expanded_strides = np.concatenate(expanded_strides, 1)
            outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
            outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides

            predictions = outputs[0]
            boxes = predictions[:, :4]
            scores = predictions[:, 4:5] * predictions[:, 5:]

            boxes_xyxy = np.ones_like(boxes)
            boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.
            boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.
            boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.
            boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.
            boxes_xyxy /= ratio
            dets = multiclass_nms(boxes_xyxy,
                                  scores,
                                  nms_thr=self.nms_thr,
                                  score_thr=self.score_thr)
            if dets is not None:
                pack_dets = (dets[:, :4], dets[:, 4], dets[:, 5])
                final_boxes, final_scores, final_cls_inds = pack_dets
                isscore = final_scores > 0.3
                iscat = final_cls_inds == 0
                isbbox = [i and j for (i, j) in zip(isscore, iscat)]
                final_boxes = final_boxes[isbbox]

        elif outputs.shape[-1] == 5:
            # onnx contains nms module

            pack_dets = (outputs[0, :, :4], outputs[0, :, 4])
            final_boxes, final_scores = pack_dets
            final_boxes /= ratio
            isscore = final_scores > 0.3
            isbbox = [i for i in isscore]
            final_boxes = final_boxes[isbbox]

        return final_boxes

rtmlib/tools/object_detection/yolox.py
Tau-J/rtmlib
from .rtmpose import RTMPose

__all__ = ['RTMPose']

rtmlib/tools/pose_estimation/__init__.py
Tau-J/rtmlib
from typing import Tuple

import numpy as np


def get_simcc_maximum(simcc_x: np.ndarray,
                      simcc_y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """Get maximum response location and value from simcc representations.

    Note:
        instance number: N
        num_keypoints: K
        heatmap height: H
        heatmap width: W

    Args:
        simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
        simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)

    Returns:
        tuple:
        - locs (np.ndarray): locations of maximum heatmap responses in shape
            (K, 2) or (N, K, 2)
        - vals (np.ndarray): values of maximum heatmap responses in shape
            (K,) or (N, K)
    """
    N, K, Wx = simcc_x.shape
    simcc_x = simcc_x.reshape(N * K, -1)
    simcc_y = simcc_y.reshape(N * K, -1)

    # get maximum value locations
    x_locs = np.argmax(simcc_x, axis=1)
    y_locs = np.argmax(simcc_y, axis=1)
    locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
    max_val_x = np.amax(simcc_x, axis=1)
    max_val_y = np.amax(simcc_y, axis=1)

    # get maximum value across x and y axis
    mask = max_val_x > max_val_y
    max_val_x[mask] = max_val_y[mask]
    vals = max_val_x
    locs[vals <= 0.] = -1

    # reshape
    locs = locs.reshape(N, K, 2)
    vals = vals.reshape(N, K)

    return locs, vals


# def convert_coco_to_openpose(keypoints, scores):
#     if keypoints.shape[1] == 17:
#         return convert_coco17_to_openpose18(keypoints, scores)
#     elif keypoints.shape[1] == 133:
#         return convert_coco133_to_openpose134(keypoints, scores)


def convert_coco_to_openpose(keypoints, scores):
    keypoints_info = np.concatenate((keypoints, scores[..., None]), axis=-1)

    # compute neck joint
    neck = np.mean(keypoints_info[:, [5, 6]], axis=1)

    # neck score when visualizing pred
    neck[:,
         2:3] = np.where(keypoints_info[:, 5, 2:3] > keypoints_info[:, 6, 2:3],
                         keypoints_info[:, 6, 2:3], keypoints_info[:, 5, 2:3])
    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_info[..., :2], keypoints_info[..., 2]
    return keypoints, scores

rtmlib/tools/pose_estimation/post_processings.py
Tau-J/rtmlib
from typing import Tuple

import cv2
import numpy as np


def bbox_xyxy2cs(bbox: np.ndarray,
                 padding: float = 1.) -> Tuple[np.ndarray, np.ndarray]:
    """Transform the bbox format from (x,y,w,h) into (center, scale)

    Args:
        bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted
            as (left, top, right, bottom)
        padding (float): BBox padding factor that will be multilied to scale.
            Default: 1.0

    Returns:
        tuple: A tuple containing center and scale.
        - np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or
            (n, 2)
        - np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or
            (n, 2)
    """
    # convert single bbox from (4, ) to (1, 4)
    dim = bbox.ndim
    if dim == 1:
        bbox = bbox[None, :]

    # get bbox center and scale
    x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3])
    center = np.hstack([x1 + x2, y1 + y2]) * 0.5
    scale = np.hstack([x2 - x1, y2 - y1]) * padding

    if dim == 1:
        center = center[0]
        scale = scale[0]

    return center, scale


def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray:
    """Rotate a point by an angle.

    Args:
        pt (np.ndarray): 2D point coordinates (x, y) in shape (2, )
        angle_rad (float): rotation angle in radian

    Returns:
        np.ndarray: Rotated point in shape (2, )
    """
    sn, cs = np.sin(angle_rad), np.cos(angle_rad)
    rot_mat = np.array([[cs, -sn], [sn, cs]])
    return rot_mat @ pt


def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """To calculate the affine matrix, three pairs of points are required. This
    function is used to get the 3rd point, given 2D points a & b.

    The 3rd point is defined by rotating vector `a - b` by 90 degrees
    anticlockwise, using b as the rotation center.

    Args:
        a (np.ndarray): The 1st point (x,y) in shape (2, )
        b (np.ndarray): The 2nd point (x,y) in shape (2, )

    Returns:
        np.ndarray: The 3rd point.
    """
    direction = a - b
    c = b + np.r_[-direction[1], direction[0]]
    return c


def get_warp_matrix(center: np.ndarray,
                    scale: np.ndarray,
                    rot: float,
                    output_size: Tuple[int, int],
                    shift: Tuple[float, float] = (0., 0.),
                    inv: bool = False) -> np.ndarray:
    """Calculate the affine transformation matrix that can warp the bbox area
    in the input image to the output size.

    Args:
        center (np.ndarray[2, ]): Center of the bounding box (x, y).
        scale (np.ndarray[2, ]): Scale of the bounding box
            wrt [width, height].
        rot (float): Rotation angle (degree).
        output_size (np.ndarray[2, ] | list(2,)): Size of the
            destination heatmaps.
        shift (0-100%): Shift translation ratio wrt the width/height.
            Default (0., 0.).
        inv (bool): Option to inverse the affine transform direction.
            (inv=False: src->dst or inv=True: dst->src)

    Returns:
        np.ndarray: A 2x3 transformation matrix
    """

    shift = np.array(shift)
    src_w = scale[0]
    dst_w = output_size[0]
    dst_h = output_size[1]

    # compute transformation matrix
    rot_rad = np.deg2rad(rot)
    src_dir = _rotate_point(np.array([0., src_w * -0.5]), rot_rad)
    dst_dir = np.array([0., dst_w * -0.5])

    # get four corners of the src rectangle in the original image
    src = np.zeros((3, 2), dtype=np.float32)
    src[0, :] = center + scale * shift
    src[1, :] = center + src_dir + scale * shift
    src[2, :] = _get_3rd_point(src[0, :], src[1, :])

    # get four corners of the dst rectangle in the input image
    dst = np.zeros((3, 2), dtype=np.float32)
    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]

    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
    dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])

    if inv:
        warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src))
    else:
        warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst))

    return warp_mat


def top_down_affine(input_size: dict, bbox_scale: dict, bbox_center: dict,
                    img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """Get the bbox image as the model input by affine transform.

    Args:
        input_size (dict): The input size of the model.
        bbox_scale (dict): The bbox scale of the img.
        bbox_center (dict): The bbox center of the img.
        img (np.ndarray): The original image.

    Returns:
        tuple: A tuple containing center and scale.
        - np.ndarray[float32]: img after affine transform.
        - np.ndarray[float32]: bbox scale after affine transform.
    """
    w, h = input_size
    warp_size = (int(w), int(h))

    # reshape bbox to fixed aspect ratio
    aspect_ratio = w / h
    b_w, b_h = np.hsplit(bbox_scale, [1])
    bbox_scale = np.where(b_w > b_h * aspect_ratio,
                          np.hstack([b_w, b_w / aspect_ratio]),
                          np.hstack([b_h * aspect_ratio, b_h]))

    # get the affine matrix
    center = bbox_center
    scale = bbox_scale
    rot = 0
    warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h))

    # do affine transform
    img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR)

    return img, bbox_scale

rtmlib/tools/pose_estimation/pre_processings.py
Tau-J/rtmlib
from typing import List, Tuple

import numpy as np

from ..base import BaseTool
from .post_processings import convert_coco_to_openpose, get_simcc_maximum
from .pre_processings import bbox_xyxy2cs, top_down_affine


class RTMPose(BaseTool):

    def __init__(self,
                 onnx_model: str,
                 model_input_size: tuple = (288, 384),
                 mean: tuple = (123.675, 116.28, 103.53),
                 std: tuple = (58.395, 57.12, 57.375),
                 to_openpose: bool = False,
                 backend: str = 'onnxruntime',
                 device: str = 'cpu'):
        super().__init__(onnx_model, model_input_size, mean, std, backend,
                         device)
        self.to_openpose = to_openpose

    def __call__(self, image: np.ndarray, bboxes: list = []):
        if len(bboxes) == 0:
            bboxes = [[0, 0, image.shape[1], image.shape[0]]]

        keypoints, scores = [], []
        for bbox in bboxes:
            img, center, scale = self.preprocess(image, bbox)
            outputs = self.inference(img)
            kpts, score = self.postprocess(outputs, center, scale)

            keypoints.append(kpts)
            scores.append(score)

        keypoints = np.concatenate(keypoints, axis=0)
        scores = np.concatenate(scores, axis=0)

        if self.to_openpose:
            keypoints, scores = convert_coco_to_openpose(keypoints, scores)

        return keypoints, scores

    def preprocess(self, img: np.ndarray, bbox: list):
        """Do preprocessing for RTMPose model inference.

        Args:
            img (np.ndarray): Input image in shape.
            bbox (list):  xyxy-format bounding box of target.

        Returns:
            tuple:
            - resized_img (np.ndarray): Preprocessed image.
            - center (np.ndarray): Center of image.
            - scale (np.ndarray): Scale of image.
        """
        bbox = np.array(bbox)

        # get center and scale
        center, scale = bbox_xyxy2cs(bbox, padding=1.25)

        # do affine transformation
        resized_img, scale = top_down_affine(self.model_input_size, scale,
                                             center, img)
        # normalize image
        if self.mean is not None:
            self.mean = np.array(self.mean)
            self.std = np.array(self.std)
            resized_img = (resized_img - self.mean) / self.std

        return resized_img, center, scale

    def postprocess(
            self,
            outputs: List[np.ndarray],
            center: Tuple[int, int],
            scale: Tuple[int, int],
            simcc_split_ratio: float = 2.0) -> Tuple[np.ndarray, np.ndarray]:
        """Postprocess for RTMPose model output.

        Args:
            outputs (np.ndarray): Output of RTMPose model.
            model_input_size (tuple): RTMPose model Input image size.
            center (tuple): Center of bbox in shape (x, y).
            scale (tuple): Scale of bbox in shape (w, h).
            simcc_split_ratio (float): Split ratio of simcc.

        Returns:
            tuple:
            - keypoints (np.ndarray): Rescaled keypoints.
            - scores (np.ndarray): Model predict scores.
        """
        # decode simcc
        simcc_x, simcc_y = outputs
        locs, scores = get_simcc_maximum(simcc_x, simcc_y)
        keypoints = locs / simcc_split_ratio

        # rescale keypoints
        keypoints = keypoints / self.model_input_size * scale
        keypoints = keypoints + center - scale / 2

        return keypoints, scores

rtmlib/tools/pose_estimation/rtmpose.py
Tau-J/rtmlib
from .body import Body
from .pose_tracker import PoseTracker
from .wholebody import Wholebody

__all__ = ['Wholebody', 'Body', 'PoseTracker']

rtmlib/tools/solution/__init__.py
Tau-J/rtmlib
'''
Example:

import cv2

from rtmlib import Body, draw_skeleton

device = 'cuda'
backend = 'onnxruntime'  # opencv, onnxruntime

cap = cv2.VideoCapture('./demo.mp4')

openpose_skeleton = True  # True for openpose-style, False for mmpose-style

body = Body(to_openpose=openpose_skeleton,
                      backend=backend,
                      device=device)

frame_idx = 0

while cap.isOpened():
    success, frame = cap.read()
    frame_idx += 1

    if not success:
        break

    keypoints, scores = body(frame)

    img_show = frame.copy()

    img_show = draw_skeleton(img_show,
                             keypoints,
                             scores,
                             openpose_skeleton=openpose_skeleton,
                             kpt_thr=0.43)

    img_show = cv2.resize(img_show, (960, 540))
    cv2.imshow('img', img_show)
    cv2.waitKey(10)

'''
import numpy as np

from .. import YOLOX, RTMPose


class Body:
    MODE = {
        'performance': {
            'det':
            'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_x_8xb8-300e_humanart-a39d44ed.zip',  # noqa
            'det_input_size': (640, 640),
            'pose':
            'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/rtmpose-x_simcc-body7_pt-body7_700e-384x288-71d7b7e9_20230629.zip',  # noqa
            'pose_input_size': (288, 384),
        },
        'lightweight': {
            'det':
            'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_tiny_8xb8-300e_humanart-6f3252f9.zip',  # noqa
            'det_input_size': (416, 416),
            'pose':
            'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/rtmpose-s_simcc-body7_pt-body7_420e-256x192-acd4a1ef_20230504.zip',  # noqa
            'pose_input_size': (192, 256),
        },
        'balanced': {
            'det':
            'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_m_8xb8-300e_humanart-c2c7a14a.zip',  # noqa
            'det_input_size': (640, 640),
            'pose':
            'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/rtmpose-m_simcc-body7_pt-body7_420e-256x192-e48f03d0_20230504.zip',  # noqa
            'pose_input_size': (192, 256),
        }
    }

    def __init__(self,
                 det: str = None,
                 det_input_size: tuple = (640, 640),
                 pose: str = None,
                 pose_input_size: tuple = (288, 384),
                 mode: str = 'balanced',
                 to_openpose: bool = False,
                 backend: str = 'onnxruntime',
                 device: str = 'cpu'):

        if det is None:
            det = self.MODE[mode]['det']
            det_input_size = self.MODE[mode]['det_input_size']

        if pose is None:
            pose = self.MODE[mode]['pose']
            pose_input_size = self.MODE[mode]['pose_input_size']

        self.det_model = YOLOX(det,
                               model_input_size=det_input_size,
                               backend=backend,
                               device=device)
        self.pose_model = RTMPose(pose,
                                  model_input_size=pose_input_size,
                                  to_openpose=to_openpose,
                                  backend=backend,
                                  device=device)

    def __call__(self, image: np.ndarray):
        bboxes = self.det_model(image)
        keypoints, scores = self.pose_model(image, bboxes=bboxes)

        return keypoints, scores

rtmlib/tools/solution/body.py
Tau-J/rtmlib
'''
Example:

import cv2
from rtmlib import PoseTracker, Wholebody, draw_skeleton

device = 'cuda'
backend = 'onnxruntime'  # opencv, onnxruntime

openpose_skeleton = True  # True for openpose-style, False for mmpose-style

cap = cv2.VideoCapture('./demo.mp4')

wholebody = PoseTracker(Wholebody,
                        det_frequency=10,  # detect every 10 frames
                        to_openpose=openpose_skeleton,
                        backend=backend, device=device)

                        frame_idx = 0

while cap.isOpened():
    success, frame = cap.read()
    frame_idx += 1

    if not success:
        break

    keypoints, scores = wholebody(frame)

    img_show = frame.copy()

    img_show = draw_skeleton(img_show,
                             keypoints,
                             scores,
                             openpose_skeleton=openpose_skeleton,
                             kpt_thr=0.43)

    img_show = cv2.resize(img_show, (960, 540))
    cv2.imshow('img', img_show)
    cv2.waitKey(10)
'''
import numpy as np


def pose_to_bbox(keypoints: np.ndarray, expansion: float = 1.25) -> np.ndarray:
    """Get bounding box from keypoints.

    Args:
        keypoints (np.ndarray): Keypoints of person.
        expansion (float): Expansion ratio of bounding box.

    Returns:
        np.ndarray: Bounding box of person.
    """
    x = keypoints[:, 0]
    y = keypoints[:, 1]
    bbox = np.array([x.min(), y.min(), x.max(), y.max()])
    center = np.array([bbox[0] + bbox[2], bbox[1] + bbox[3]]) / 2
    bbox = np.concatenate([
        center - (center - bbox[:2]) * expansion,
        center + (bbox[2:] - center) * expansion
    ])
    return bbox


class PoseTracker:
    """Pose tracker for wholebody pose estimation.

    Args:
        solution (type): rtmlib solutions, e.g. Wholebody, Body, etc.
        det_frequency (int): Frequency of object detection.
        mode (str): 'performance', 'lightweight', or 'balanced'.
        to_openpose (bool): Whether to use openpose-style skeleton.
        backend (str): Backend of pose estimation model.
        device (str): Device of pose estimation model.
    """

    def __init__(self,
                 solution: type,
                 det_frequency: int = 1,
                 mode: str = 'balanced',
                 to_openpose: bool = False,
                 backend: str = 'onnxruntime',
                 device: str = 'cpu'):

        model = solution(mode=mode,
                         to_openpose=to_openpose,
                         backend=backend,
                         device=device)

        self.det_model = model.det_model
        self.pose_model = model.pose_model

        self.det_frequency = det_frequency
        self.reset()

    def reset(self):
        """Reset pose tracker."""
        self.cnt = 0
        self.instance_list = []

    def __call__(self, image: np.ndarray):

        if self.cnt % self.det_frequency == 0:
            bboxes = self.det_model(image)
        else:
            bboxes = self.instance_list

        keypoints, scores = self.pose_model(image, bboxes=bboxes)

        instances = []
        for kpts in keypoints:
            instances.append(pose_to_bbox(kpts))

        self.instance_list = instances
        self.cnt += 1

        return keypoints, scores

rtmlib/tools/solution/pose_tracker.py
Tau-J/rtmlib
# poted from controlnet repo
from typing import List, NamedTuple, Optional


class Keypoint(NamedTuple):
    x: float
    y: float
    score: float = 1.0
    id: int = -1


class BodyResult(NamedTuple):
    # Note: Using `Optional` instead of `|` operator as the ladder is a Python
    # 3.10 feature.
    # Annotator code should be Python 3.8 Compatible, as controlnet repo uses
    # Python 3.8 environment.
    # https://github.com/lllyasviel/ControlNet/blob/d3284fcd0972c510635a4f5abe2eeb71dc0de524/environment.yaml#L6
    keypoints: List[Optional[Keypoint]]
    total_score: float = 0.0
    total_parts: int = 0


HandResult = List[Keypoint]
FaceResult = List[Keypoint]


class PoseResult(NamedTuple):
    body: BodyResult
    left_hand: Optional[HandResult]
    right_hand: Optional[HandResult]
    face: Optional[FaceResult]

rtmlib/tools/solution/utils/types.py
Tau-J/rtmlib
'''
Example:

import cv2

from rtmlib import Wholebody, draw_skeleton

device = 'cuda'
backend = 'onnxruntime'  # opencv, onnxruntime

cap = cv2.VideoCapture('./demo.mp4')

openpose_skeleton = True  # True for openpose-style, False for mmpose-style

wholebody = Wholebody(to_openpose=openpose_skeleton,
                      backend=backend,
                      device=device)

frame_idx = 0

while cap.isOpened():
    success, frame = cap.read()
    frame_idx += 1

    if not success:
        break

    keypoints, scores = wholebody(frame)

    img_show = frame.copy()

    img_show = draw_skeleton(img_show,
                             keypoints,
                             scores,
                             openpose_skeleton=openpose_skeleton,
                             kpt_thr=0.43)

    img_show = cv2.resize(img_show, (960, 540))
    cv2.imshow('img', img_show)
    cv2.waitKey(10)

'''
from typing import List, Optional

import numpy as np

from .. import YOLOX, RTMPose
from .utils.types import BodyResult, Keypoint, PoseResult


class Wholebody:

    MODE = {
        'performance': {
            'det':
            'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_x_8xb8-300e_humanart-a39d44ed.zip',  # noqa
            'det_input_size': (640, 640),
            'pose':
            'https://download.openmmlab.com/mmpose/v1/projects/rtmw/onnx_sdk/rtmw-x_simcc-cocktail13_pt-ucoco_270e-384x288-0949e3a9_20230925.zip',  # noqa
            'pose_input_size': (288, 384),
        },
        'lightweight': {
            'det':
            'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_tiny_8xb8-300e_humanart-6f3252f9.zip',  # noqa
            'det_input_size': (416, 416),
            'pose':
            'https://download.openmmlab.com/mmpose/v1/projects/rtmw/onnx_sdk/rtmw-x_simcc-cocktail13_pt-ucoco_270e-256x192-fbef0d61_20230925.zip',  # noqa
            'pose_input_size': (192, 256),
        },
        'balanced': {
            'det':
            'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_m_8xb8-300e_humanart-c2c7a14a.zip',  # noqa
            'det_input_size': (640, 640),
            'pose':
            'https://download.openmmlab.com/mmpose/v1/projects/rtmw/onnx_sdk/rtmw-x_simcc-cocktail13_pt-ucoco_270e-256x192-fbef0d61_20230925.zip',  # noqa
            'pose_input_size': (192, 256),
        }
    }

    def __init__(self,
                 det: str = None,
                 det_input_size: tuple = (640, 640),
                 pose: str = None,
                 pose_input_size: tuple = (288, 384),
                 mode: str = 'balanced',
                 to_openpose: bool = False,
                 backend: str = 'onnxruntime',
                 device: str = 'cpu'):

        if det is None:
            det = self.MODE[mode]['det']
            det_input_size = self.MODE[mode]['det_input_size']

        if pose is None:
            pose = self.MODE[mode]['pose']
            pose_input_size = self.MODE[mode]['pose_input_size']

        self.det_model = YOLOX(det,
                               model_input_size=det_input_size,
                               backend=backend,
                               device=device)
        self.pose_model = RTMPose(pose,
                                  model_input_size=pose_input_size,
                                  to_openpose=to_openpose,
                                  backend=backend,
                                  device=device)

    def __call__(self, image: np.ndarray):
        bboxes = self.det_model(image)
        keypoints, scores = self.pose_model(image, bboxes=bboxes)

        return keypoints, scores

    @staticmethod
    def format_result(keypoints_info: np.ndarray) -> List[PoseResult]:

        def format_keypoint_part(
                part: np.ndarray) -> Optional[List[Optional[Keypoint]]]:
            keypoints = [
                Keypoint(x, y, score, i) if score >= 0.3 else None
                for i, (x, y, score) in enumerate(part)
            ]
            return (None if all(keypoint is None
                                for keypoint in keypoints) else keypoints)

        def total_score(
                keypoints: Optional[List[Optional[Keypoint]]]) -> float:
            return (sum(
                keypoint.score for keypoint in keypoints
                if keypoint is not None) if keypoints is not None else 0.0)

        pose_results = []

        for instance in keypoints_info:
            body_keypoints = format_keypoint_part(
                instance[:18]) or ([None] * 18)
            left_hand = format_keypoint_part(instance[92:113])
            right_hand = format_keypoint_part(instance[113:134])
            face = format_keypoint_part(instance[24:92])

            # Openpose face consists of 70 points in total, while RTMPose only
            # provides 68 points. Padding the last 2 points.
            if face is not None:
                # left eye
                face.append(body_keypoints[14])
                # right eye
                face.append(body_keypoints[15])

            body = BodyResult(body_keypoints, total_score(body_keypoints),
                              len(body_keypoints))
            pose_results.append(PoseResult(body, left_hand, right_hand, face))

        return pose_results

rtmlib/tools/solution/wholebody.py
Tau-J/rtmlib
__version__ = '0.0.6'
short_version = __version__


def parse_version_info(version_str):
    """Parse a version string into a tuple.

    Args:
        version_str (str): The version string.
    Returns:
        tuple[int | str]: The version info, e.g., "1.3.0" is parsed into
            (1, 3, 0), and "2.0.0rc1" is parsed into (2, 0, 0, 'rc1').
    """
    version_info = []
    for x in version_str.split('.'):
        if x.isdigit():
            version_info.append(int(x))
        elif x.find('rc') != -1:
            patch_version = x.split('rc')
            version_info.append(int(patch_version[0]))
            version_info.append(f'rc{patch_version[1]}')
        elif x.find('b') != -1:
            patch_version = x.split('b')
            version_info.append(int(patch_version[0]))
            version_info.append(f'b{patch_version[1]}')
    return tuple(version_info)


version_info = parse_version_info(__version__)

rtmlib/version.py
Tau-J/rtmlib
from .coco17 import coco17
from .coco133 import coco133
from .openpose18 import openpose18
from .openpose134 import openpose134

__all__ = ['coco17', 'openpose18', 'coco133', 'openpose134']

rtmlib/visualization/skeleton/__init__.py
Tau-J/rtmlib
coco133 = dict(
    name='coco133',
    keypoint_info={
        0:
        dict(name='nose', id=0, color=[51, 153, 255], swap=''),
        1:
        dict(name='left_eye', id=1, color=[51, 153, 255], swap='right_eye'),
        2:
        dict(name='right_eye', id=2, color=[51, 153, 255], swap='left_eye'),
        3:
        dict(name='left_ear', id=3, color=[51, 153, 255], swap='right_ear'),
        4:
        dict(name='right_ear', id=4, color=[51, 153, 255], swap='left_ear'),
        5:
        dict(name='left_shoulder',
             id=5,
             color=[0, 255, 0],
             swap='right_shoulder'),
        6:
        dict(name='right_shoulder',
             id=6,
             color=[255, 128, 0],
             swap='left_shoulder'),
        7:
        dict(name='left_elbow', id=7, color=[0, 255, 0], swap='right_elbow'),
        8:
        dict(name='right_elbow', id=8, color=[255, 128, 0], swap='left_elbow'),
        9:
        dict(name='left_wrist', id=9, color=[0, 255, 0], swap='right_wrist'),
        10:
        dict(name='right_wrist', id=10, color=[255, 128, 0],
             swap='left_wrist'),
        11:
        dict(name='left_hip', id=11, color=[0, 255, 0], swap='right_hip'),
        12:
        dict(name='right_hip', id=12, color=[255, 128, 0], swap='left_hip'),
        13:
        dict(name='left_knee', id=13, color=[0, 255, 0], swap='right_knee'),
        14:
        dict(name='right_knee', id=14, color=[255, 128, 0], swap='left_knee'),
        15:
        dict(name='left_ankle', id=15, color=[0, 255, 0], swap='right_ankle'),
        16:
        dict(name='right_ankle', id=16, color=[255, 128, 0],
             swap='left_ankle'),
        17:
        dict(name='left_big_toe',
             id=17,
             color=[255, 128, 0],
             swap='right_big_toe'),
        18:
        dict(name='left_small_toe',
             id=18,
             color=[255, 128, 0],
             swap='right_small_toe'),
        19:
        dict(name='left_heel', id=19, color=[255, 128, 0], swap='right_heel'),
        20:
        dict(name='right_big_toe',
             id=20,
             color=[255, 128, 0],
             swap='left_big_toe'),
        21:
        dict(name='right_small_toe',
             id=21,
             color=[255, 128, 0],
             swap='left_small_toe'),
        22:
        dict(name='right_heel', id=22, color=[255, 128, 0], swap='left_heel'),
        23:
        dict(name='face-0', id=23, color=[255, 255, 255], swap='face-16'),
        24:
        dict(name='face-1', id=24, color=[255, 255, 255], swap='face-15'),
        25:
        dict(name='face-2', id=25, color=[255, 255, 255], swap='face-14'),
        26:
        dict(name='face-3', id=26, color=[255, 255, 255], swap='face-13'),
        27:
        dict(name='face-4', id=27, color=[255, 255, 255], swap='face-12'),
        28:
        dict(name='face-5', id=28, color=[255, 255, 255], swap='face-11'),
        29:
        dict(name='face-6', id=29, color=[255, 255, 255], swap='face-10'),
        30:
        dict(name='face-7', id=30, color=[255, 255, 255], swap='face-9'),
        31:
        dict(name='face-8', id=31, color=[255, 255, 255], swap=''),
        32:
        dict(name='face-9', id=32, color=[255, 255, 255], swap='face-7'),
        33:
        dict(name='face-10', id=33, color=[255, 255, 255], swap='face-6'),
        34:
        dict(name='face-11', id=34, color=[255, 255, 255], swap='face-5'),
        35:
        dict(name='face-12', id=35, color=[255, 255, 255], swap='face-4'),
        36:
        dict(name='face-13', id=36, color=[255, 255, 255], swap='face-3'),
        37:
        dict(name='face-14', id=37, color=[255, 255, 255], swap='face-2'),
        38:
        dict(name='face-15', id=38, color=[255, 255, 255], swap='face-1'),
        39:
        dict(name='face-16', id=39, color=[255, 255, 255], swap='face-0'),
        40:
        dict(name='face-17', id=40, color=[255, 255, 255], swap='face-26'),
        41:
        dict(name='face-18', id=41, color=[255, 255, 255], swap='face-25'),
        42:
        dict(name='face-19', id=42, color=[255, 255, 255], swap='face-24'),
        43:
        dict(name='face-20', id=43, color=[255, 255, 255], swap='face-23'),
        44:
        dict(name='face-21', id=44, color=[255, 255, 255], swap='face-22'),
        45:
        dict(name='face-22', id=45, color=[255, 255, 255], swap='face-21'),
        46:
        dict(name='face-23', id=46, color=[255, 255, 255], swap='face-20'),
        47:
        dict(name='face-24', id=47, color=[255, 255, 255], swap='face-19'),
        48:
        dict(name='face-25', id=48, color=[255, 255, 255], swap='face-18'),
        49:
        dict(name='face-26', id=49, color=[255, 255, 255], swap='face-17'),
        50:
        dict(name='face-27', id=50, color=[255, 255, 255], swap=''),
        51:
        dict(name='face-28', id=51, color=[255, 255, 255], swap=''),
        52:
        dict(name='face-29', id=52, color=[255, 255, 255], swap=''),
        53:
        dict(name='face-30', id=53, color=[255, 255, 255], swap=''),
        54:
        dict(name='face-31', id=54, color=[255, 255, 255], swap='face-35'),
        55:
        dict(name='face-32', id=55, color=[255, 255, 255], swap='face-34'),
        56:
        dict(name='face-33', id=56, color=[255, 255, 255], swap=''),
        57:
        dict(name='face-34', id=57, color=[255, 255, 255], swap='face-32'),
        58:
        dict(name='face-35', id=58, color=[255, 255, 255], swap='face-31'),
        59:
        dict(name='face-36', id=59, color=[255, 255, 255], swap='face-45'),
        60:
        dict(name='face-37', id=60, color=[255, 255, 255], swap='face-44'),
        61:
        dict(name='face-38', id=61, color=[255, 255, 255], swap='face-43'),
        62:
        dict(name='face-39', id=62, color=[255, 255, 255], swap='face-42'),
        63:
        dict(name='face-40', id=63, color=[255, 255, 255], swap='face-47'),
        64:
        dict(name='face-41', id=64, color=[255, 255, 255], swap='face-46'),
        65:
        dict(name='face-42', id=65, color=[255, 255, 255], swap='face-39'),
        66:
        dict(name='face-43', id=66, color=[255, 255, 255], swap='face-38'),
        67:
        dict(name='face-44', id=67, color=[255, 255, 255], swap='face-37'),
        68:
        dict(name='face-45', id=68, color=[255, 255, 255], swap='face-36'),
        69:
        dict(name='face-46', id=69, color=[255, 255, 255], swap='face-41'),
        70:
        dict(name='face-47', id=70, color=[255, 255, 255], swap='face-40'),
        71:
        dict(name='face-48', id=71, color=[255, 255, 255], swap='face-54'),
        72:
        dict(name='face-49', id=72, color=[255, 255, 255], swap='face-53'),
        73:
        dict(name='face-50', id=73, color=[255, 255, 255], swap='face-52'),
        74:
        dict(name='face-51', id=74, color=[255, 255, 255], swap=''),
        75:
        dict(name='face-52', id=75, color=[255, 255, 255], swap='face-50'),
        76:
        dict(name='face-53', id=76, color=[255, 255, 255], swap='face-49'),
        77:
        dict(name='face-54', id=77, color=[255, 255, 255], swap='face-48'),
        78:
        dict(name='face-55', id=78, color=[255, 255, 255], swap='face-59'),
        79:
        dict(name='face-56', id=79, color=[255, 255, 255], swap='face-58'),
        80:
        dict(name='face-57', id=80, color=[255, 255, 255], swap=''),
        81:
        dict(name='face-58', id=81, color=[255, 255, 255], swap='face-56'),
        82:
        dict(name='face-59', id=82, color=[255, 255, 255], swap='face-55'),
        83:
        dict(name='face-60', id=83, color=[255, 255, 255], swap='face-64'),
        84:
        dict(name='face-61', id=84, color=[255, 255, 255], swap='face-63'),
        85:
        dict(name='face-62', id=85, color=[255, 255, 255], swap=''),
        86:
        dict(name='face-63', id=86, color=[255, 255, 255], swap='face-61'),
        87:
        dict(name='face-64', id=87, color=[255, 255, 255], swap='face-60'),
        88:
        dict(name='face-65', id=88, color=[255, 255, 255], swap='face-67'),
        89:
        dict(name='face-66', id=89, color=[255, 255, 255], swap=''),
        90:
        dict(name='face-67', id=90, color=[255, 255, 255], swap='face-65'),
        91:
        dict(name='left_hand_root',
             id=91,
             color=[255, 255, 255],
             swap='right_hand_root'),
        92:
        dict(name='left_thumb1',
             id=92,
             color=[255, 128, 0],
             swap='right_thumb1'),
        93:
        dict(name='left_thumb2',
             id=93,
             color=[255, 128, 0],
             swap='right_thumb2'),
        94:
        dict(name='left_thumb3',
             id=94,
             color=[255, 128, 0],
             swap='right_thumb3'),
        95:
        dict(name='left_thumb4',
             id=95,
             color=[255, 128, 0],
             swap='right_thumb4'),
        96:
        dict(name='left_forefinger1',
             id=96,
             color=[255, 153, 255],
             swap='right_forefinger1'),
        97:
        dict(name='left_forefinger2',
             id=97,
             color=[255, 153, 255],
             swap='right_forefinger2'),
        98:
        dict(name='left_forefinger3',
             id=98,
             color=[255, 153, 255],
             swap='right_forefinger3'),
        99:
        dict(name='left_forefinger4',
             id=99,
             color=[255, 153, 255],
             swap='right_forefinger4'),
        100:
        dict(name='left_middle_finger1',
             id=100,
             color=[102, 178, 255],
             swap='right_middle_finger1'),
        101:
        dict(name='left_middle_finger2',
             id=101,
             color=[102, 178, 255],
             swap='right_middle_finger2'),
        102:
        dict(name='left_middle_finger3',
             id=102,
             color=[102, 178, 255],
             swap='right_middle_finger3'),
        103:
        dict(name='left_middle_finger4',
             id=103,
             color=[102, 178, 255],
             swap='right_middle_finger4'),
        104:
        dict(name='left_ring_finger1',
             id=104,
             color=[255, 51, 51],
             swap='right_ring_finger1'),
        105:
        dict(name='left_ring_finger2',
             id=105,
             color=[255, 51, 51],
             swap='right_ring_finger2'),
        106:
        dict(name='left_ring_finger3',
             id=106,
             color=[255, 51, 51],
             swap='right_ring_finger3'),
        107:
        dict(name='left_ring_finger4',
             id=107,
             color=[255, 51, 51],
             swap='right_ring_finger4'),
        108:
        dict(name='left_pinky_finger1',
             id=108,
             color=[0, 255, 0],
             swap='right_pinky_finger1'),
        109:
        dict(name='left_pinky_finger2',
             id=109,
             color=[0, 255, 0],
             swap='right_pinky_finger2'),
        110:
        dict(name='left_pinky_finger3',
             id=110,
             color=[0, 255, 0],
             swap='right_pinky_finger3'),
        111:
        dict(name='left_pinky_finger4',
             id=111,
             color=[0, 255, 0],
             swap='right_pinky_finger4'),
        112:
        dict(name='right_hand_root',
             id=112,
             color=[255, 255, 255],
             swap='left_hand_root'),
        113:
        dict(name='right_thumb1',
             id=113,
             color=[255, 128, 0],
             swap='left_thumb1'),
        114:
        dict(name='right_thumb2',
             id=114,
             color=[255, 128, 0],
             swap='left_thumb2'),
        115:
        dict(name='right_thumb3',
             id=115,
             color=[255, 128, 0],
             swap='left_thumb3'),
        116:
        dict(name='right_thumb4',
             id=116,
             color=[255, 128, 0],
             swap='left_thumb4'),
        117:
        dict(name='right_forefinger1',
             id=117,
             color=[255, 153, 255],
             swap='left_forefinger1'),
        118:
        dict(name='right_forefinger2',
             id=118,
             color=[255, 153, 255],
             swap='left_forefinger2'),
        119:
        dict(name='right_forefinger3',
             id=119,
             color=[255, 153, 255],
             swap='left_forefinger3'),
        120:
        dict(name='right_forefinger4',
             id=120,
             color=[255, 153, 255],
             swap='left_forefinger4'),
        121:
        dict(name='right_middle_finger1',
             id=121,
             color=[102, 178, 255],
             swap='left_middle_finger1'),
        122:
        dict(name='right_middle_finger2',
             id=122,
             color=[102, 178, 255],
             swap='left_middle_finger2'),
        123:
        dict(name='right_middle_finger3',
             id=123,
             color=[102, 178, 255],
             swap='left_middle_finger3'),
        124:
        dict(name='right_middle_finger4',
             id=124,
             color=[102, 178, 255],
             swap='left_middle_finger4'),
        125:
        dict(name='right_ring_finger1',
             id=125,
             color=[255, 51, 51],
             swap='left_ring_finger1'),
        126:
        dict(name='right_ring_finger2',
             id=126,
             color=[255, 51, 51],
             swap='left_ring_finger2'),
        127:
        dict(name='right_ring_finger3',
             id=127,
             color=[255, 51, 51],
             swap='left_ring_finger3'),
        128:
        dict(name='right_ring_finger4',
             id=128,
             color=[255, 51, 51],
             swap='left_ring_finger4'),
        129:
        dict(name='right_pinky_finger1',
             id=129,
             color=[0, 255, 0],
             swap='left_pinky_finger1'),
        130:
        dict(name='right_pinky_finger2',
             id=130,
             color=[0, 255, 0],
             swap='left_pinky_finger2'),
        131:
        dict(name='right_pinky_finger3',
             id=131,
             color=[0, 255, 0],
             swap='left_pinky_finger3'),
        132:
        dict(name='right_pinky_finger4',
             id=132,
             color=[0, 255, 0],
             swap='left_pinky_finger4')
    },
    skeleton_info={
        0:
        dict(link=('left_ankle', 'left_knee'), id=0, color=[0, 255, 0]),
        1:
        dict(link=('left_knee', 'left_hip'), id=1, color=[0, 255, 0]),
        2:
        dict(link=('right_ankle', 'right_knee'), id=2, color=[255, 128, 0]),
        3:
        dict(link=('right_knee', 'right_hip'), id=3, color=[255, 128, 0]),
        4:
        dict(link=('left_hip', 'right_hip'), id=4, color=[51, 153, 255]),
        5:
        dict(link=('left_shoulder', 'left_hip'), id=5, color=[51, 153, 255]),
        6:
        dict(link=('right_shoulder', 'right_hip'), id=6, color=[51, 153, 255]),
        7:
        dict(link=('left_shoulder', 'right_shoulder'),
             id=7,
             color=[51, 153, 255]),
        8:
        dict(link=('left_shoulder', 'left_elbow'), id=8, color=[0, 255, 0]),
        9:
        dict(link=('right_shoulder', 'right_elbow'), id=9, color=[255, 128,
                                                                  0]),
        10:
        dict(link=('left_elbow', 'left_wrist'), id=10, color=[0, 255, 0]),
        11:
        dict(link=('right_elbow', 'right_wrist'), id=11, color=[255, 128, 0]),
        12:
        dict(link=('left_eye', 'right_eye'), id=12, color=[51, 153, 255]),
        13:
        dict(link=('nose', 'left_eye'), id=13, color=[51, 153, 255]),
        14:
        dict(link=('nose', 'right_eye'), id=14, color=[51, 153, 255]),
        15:
        dict(link=('left_eye', 'left_ear'), id=15, color=[51, 153, 255]),
        16:
        dict(link=('right_eye', 'right_ear'), id=16, color=[51, 153, 255]),
        17:
        dict(link=('left_ear', 'left_shoulder'), id=17, color=[51, 153, 255]),
        18:
        dict(link=('right_ear', 'right_shoulder'), id=18, color=[51, 153,
                                                                 255]),
        19:
        dict(link=('left_ankle', 'left_big_toe'), id=19, color=[0, 255, 0]),
        20:
        dict(link=('left_ankle', 'left_small_toe'), id=20, color=[0, 255, 0]),
        21:
        dict(link=('left_ankle', 'left_heel'), id=21, color=[0, 255, 0]),
        22:
        dict(link=('right_ankle', 'right_big_toe'), id=22, color=[255, 128,
                                                                  0]),
        23:
        dict(link=('right_ankle', 'right_small_toe'),
             id=23,
             color=[255, 128, 0]),
        24:
        dict(link=('right_ankle', 'right_heel'), id=24, color=[255, 128, 0]),
        25:
        dict(link=('left_hand_root', 'left_thumb1'),
             id=25,
             color=[255, 128, 0]),
        26:
        dict(link=('left_thumb1', 'left_thumb2'), id=26, color=[255, 128, 0]),
        27:
        dict(link=('left_thumb2', 'left_thumb3'), id=27, color=[255, 128, 0]),
        28:
        dict(link=('left_thumb3', 'left_thumb4'), id=28, color=[255, 128, 0]),
        29:
        dict(link=('left_hand_root', 'left_forefinger1'),
             id=29,
             color=[255, 153, 255]),
        30:
        dict(link=('left_forefinger1', 'left_forefinger2'),
             id=30,
             color=[255, 153, 255]),
        31:
        dict(link=('left_forefinger2', 'left_forefinger3'),
             id=31,
             color=[255, 153, 255]),
        32:
        dict(link=('left_forefinger3', 'left_forefinger4'),
             id=32,
             color=[255, 153, 255]),
        33:
        dict(link=('left_hand_root', 'left_middle_finger1'),
             id=33,
             color=[102, 178, 255]),
        34:
        dict(link=('left_middle_finger1', 'left_middle_finger2'),
             id=34,
             color=[102, 178, 255]),
        35:
        dict(link=('left_middle_finger2', 'left_middle_finger3'),
             id=35,
             color=[102, 178, 255]),
        36:
        dict(link=('left_middle_finger3', 'left_middle_finger4'),
             id=36,
             color=[102, 178, 255]),
        37:
        dict(link=('left_hand_root', 'left_ring_finger1'),
             id=37,
             color=[255, 51, 51]),
        38:
        dict(link=('left_ring_finger1', 'left_ring_finger2'),
             id=38,
             color=[255, 51, 51]),
        39:
        dict(link=('left_ring_finger2', 'left_ring_finger3'),
             id=39,
             color=[255, 51, 51]),
        40:
        dict(link=('left_ring_finger3', 'left_ring_finger4'),
             id=40,
             color=[255, 51, 51]),
        41:
        dict(link=('left_hand_root', 'left_pinky_finger1'),
             id=41,
             color=[0, 255, 0]),
        42:
        dict(link=('left_pinky_finger1', 'left_pinky_finger2'),
             id=42,
             color=[0, 255, 0]),
        43:
        dict(link=('left_pinky_finger2', 'left_pinky_finger3'),
             id=43,
             color=[0, 255, 0]),
        44:
        dict(link=('left_pinky_finger3', 'left_pinky_finger4'),
             id=44,
             color=[0, 255, 0]),
        45:
        dict(link=('right_hand_root', 'right_thumb1'),
             id=45,
             color=[255, 128, 0]),
        46:
        dict(link=('right_thumb1', 'right_thumb2'), id=46, color=[255, 128,
                                                                  0]),
        47:
        dict(link=('right_thumb2', 'right_thumb3'), id=47, color=[255, 128,
                                                                  0]),
        48:
        dict(link=('right_thumb3', 'right_thumb4'), id=48, color=[255, 128,
                                                                  0]),
        49:
        dict(link=('right_hand_root', 'right_forefinger1'),
             id=49,
             color=[255, 153, 255]),
        50:
        dict(link=('right_forefinger1', 'right_forefinger2'),
             id=50,
             color=[255, 153, 255]),
        51:
        dict(link=('right_forefinger2', 'right_forefinger3'),
             id=51,
             color=[255, 153, 255]),
        52:
        dict(link=('right_forefinger3', 'right_forefinger4'),
             id=52,
             color=[255, 153, 255]),
        53:
        dict(link=('right_hand_root', 'right_middle_finger1'),
             id=53,
             color=[102, 178, 255]),
        54:
        dict(link=('right_middle_finger1', 'right_middle_finger2'),
             id=54,
             color=[102, 178, 255]),
        55:
        dict(link=('right_middle_finger2', 'right_middle_finger3'),
             id=55,
             color=[102, 178, 255]),
        56:
        dict(link=('right_middle_finger3', 'right_middle_finger4'),
             id=56,
             color=[102, 178, 255]),
        57:
        dict(link=('right_hand_root', 'right_ring_finger1'),
             id=57,
             color=[255, 51, 51]),
        58:
        dict(link=('right_ring_finger1', 'right_ring_finger2'),
             id=58,
             color=[255, 51, 51]),
        59:
        dict(link=('right_ring_finger2', 'right_ring_finger3'),
             id=59,
             color=[255, 51, 51]),
        60:
        dict(link=('right_ring_finger3', 'right_ring_finger4'),
             id=60,
             color=[255, 51, 51]),
        61:
        dict(link=('right_hand_root', 'right_pinky_finger1'),
             id=61,
             color=[0, 255, 0]),
        62:
        dict(link=('right_pinky_finger1', 'right_pinky_finger2'),
             id=62,
             color=[0, 255, 0]),
        63:
        dict(link=('right_pinky_finger2', 'right_pinky_finger3'),
             id=63,
             color=[0, 255, 0]),
        64:
        dict(link=('right_pinky_finger3', 'right_pinky_finger4'),
             id=64,
             color=[0, 255, 0])
    },
)

rtmlib/visualization/skeleton/coco133.py
Tau-J/rtmlib
coco17 = dict(name='coco17',
              keypoint_info={
                  0:
                  dict(name='nose', id=0, color=[51, 153, 255], swap=''),
                  1:
                  dict(name='left_eye',
                       id=1,
                       color=[51, 153, 255],
                       swap='right_eye'),
                  2:
                  dict(name='right_eye',
                       id=2,
                       color=[51, 153, 255],
                       swap='left_eye'),
                  3:
                  dict(name='left_ear',
                       id=3,
                       color=[51, 153, 255],
                       swap='right_ear'),
                  4:
                  dict(name='right_ear',
                       id=4,
                       color=[51, 153, 255],
                       swap='left_ear'),
                  5:
                  dict(name='left_shoulder',
                       id=5,
                       color=[0, 255, 0],
                       swap='right_shoulder'),
                  6:
                  dict(name='right_shoulder',
                       id=6,
                       color=[255, 128, 0],
                       swap='left_shoulder'),
                  7:
                  dict(name='left_elbow',
                       id=7,
                       color=[0, 255, 0],
                       swap='right_elbow'),
                  8:
                  dict(name='right_elbow',
                       id=8,
                       color=[255, 128, 0],
                       swap='left_elbow'),
                  9:
                  dict(name='left_wrist',
                       id=9,
                       color=[0, 255, 0],
                       swap='right_wrist'),
                  10:
                  dict(name='right_wrist',
                       id=10,
                       color=[255, 128, 0],
                       swap='left_wrist'),
                  11:
                  dict(name='left_hip',
                       id=11,
                       color=[0, 255, 0],
                       swap='right_hip'),
                  12:
                  dict(name='right_hip',
                       id=12,
                       color=[255, 128, 0],
                       swap='left_hip'),
                  13:
                  dict(name='left_knee',
                       id=13,
                       color=[0, 255, 0],
                       swap='right_knee'),
                  14:
                  dict(name='right_knee',
                       id=14,
                       color=[255, 128, 0],
                       swap='left_knee'),
                  15:
                  dict(name='left_ankle',
                       id=15,
                       color=[0, 255, 0],
                       swap='right_ankle'),
                  16:
                  dict(name='right_ankle',
                       id=16,
                       color=[255, 128, 0],
                       swap='left_ankle')
              },
              skeleton_info={
                  0:
                  dict(link=('left_ankle', 'left_knee'),
                       id=0,
                       color=[0, 255, 0]),
                  1:
                  dict(link=('left_knee', 'left_hip'), id=1, color=[0, 255,
                                                                    0]),
                  2:
                  dict(link=('right_ankle', 'right_knee'),
                       id=2,
                       color=[255, 128, 0]),
                  3:
                  dict(link=('right_knee', 'right_hip'),
                       id=3,
                       color=[255, 128, 0]),
                  4:
                  dict(link=('left_hip', 'right_hip'),
                       id=4,
                       color=[51, 153, 255]),
                  5:
                  dict(link=('left_shoulder', 'left_hip'),
                       id=5,
                       color=[51, 153, 255]),
                  6:
                  dict(link=('right_shoulder', 'right_hip'),
                       id=6,
                       color=[51, 153, 255]),
                  7:
                  dict(link=('left_shoulder', 'right_shoulder'),
                       id=7,
                       color=[51, 153, 255]),
                  8:
                  dict(link=('left_shoulder', 'left_elbow'),
                       id=8,
                       color=[0, 255, 0]),
                  9:
                  dict(link=('right_shoulder', 'right_elbow'),
                       id=9,
                       color=[255, 128, 0]),
                  10:
                  dict(link=('left_elbow', 'left_wrist'),
                       id=10,
                       color=[0, 255, 0]),
                  11:
                  dict(link=('right_elbow', 'right_wrist'),
                       id=11,
                       color=[255, 128, 0]),
                  12:
                  dict(link=('left_eye', 'right_eye'),
                       id=12,
                       color=[51, 153, 255]),
                  13:
                  dict(link=('nose', 'left_eye'), id=13, color=[51, 153, 255]),
                  14:
                  dict(link=('nose', 'right_eye'), id=14, color=[51, 153,
                                                                 255]),
                  15:
                  dict(link=('left_eye', 'left_ear'),
                       id=15,
                       color=[51, 153, 255]),
                  16:
                  dict(link=('right_eye', 'right_ear'),
                       id=16,
                       color=[51, 153, 255]),
                  17:
                  dict(link=('left_ear', 'left_shoulder'),
                       id=17,
                       color=[51, 153, 255]),
                  18:
                  dict(link=('right_ear', 'right_shoulder'),
                       id=18,
                       color=[51, 153, 255])
              })

rtmlib/visualization/skeleton/coco17.py
Tau-J/rtmlib
openpose134 = dict(
    name='openpose134',
    keypoint_info={
        0:
        dict(name='nose', id=0, color=[255, 0, 0], swap=''),
        1:
        dict(name='neck', id=1, color=[255, 85, 0], swap=''),
        2:
        dict(name='right_shoulder',
             id=2,
             color=[255, 170, 0],
             swap='left_shoulder'),
        3:
        dict(name='right_elbow', id=3, color=[255, 255, 0], swap='left_elbow'),
        4:
        dict(name='right_wrist', id=4, color=[170, 255, 0], swap='left_wrist'),
        5:
        dict(name='left_shoulder',
             id=5,
             color=[85, 255, 0],
             swap='right_shoulder'),
        6:
        dict(name='left_elbow', id=6, color=[0, 255, 0], swap='right_elbow'),
        7:
        dict(name='left_wrist', id=7, color=[0, 255, 85], swap='right_wrist'),
        8:
        dict(name='right_hip', id=8, color=[0, 255, 170], swap='left_hip'),
        9:
        dict(name='right_knee', id=9, color=[0, 255, 255], swap='left_knee'),
        10:
        dict(name='right_ankle', id=10, color=[0, 170, 255],
             swap='left_ankle'),
        11:
        dict(name='left_hip', id=11, color=[0, 85, 255], swap='right_hip'),
        12:
        dict(name='left_knee', id=12, color=[0, 0, 255], swap='right_knee'),
        13:
        dict(name='left_ankle', id=13, color=[85, 0, 255], swap='right_ankle'),
        14:
        dict(name='right_eye', id=14, color=[170, 0, 255], swap='left_eye'),
        15:
        dict(name='left_eye', id=15, color=[255, 0, 255], swap='right_eye'),
        16:
        dict(name='right_ear', id=16, color=[255, 0, 170], swap='left_ear'),
        17:
        dict(name='left_ear', id=17, color=[255, 0, 85], swap='right_ear'),
        18:
        dict(name='left_big_toe', id=17, color=[0, 0, 0],
             swap='right_big_toe'),
        19:
        dict(name='left_small_toe',
             id=18,
             color=[0, 0, 0],
             swap='right_small_toe'),
        20:
        dict(name='left_heel', id=19, color=[0, 0, 0], swap='right_heel'),
        21:
        dict(name='right_big_toe', id=20, color=[0, 0, 0],
             swap='left_big_toe'),
        22:
        dict(name='right_small_toe',
             id=21,
             color=[0, 0, 0],
             swap='left_small_toe'),
        23:
        dict(name='right_heel', id=22, color=[0, 0, 0], swap='left_heel'),
        24:
        dict(name='face-0', id=23, color=[255, 255, 255], swap='face-16'),
        25:
        dict(name='face-1', id=24, color=[255, 255, 255], swap='face-15'),
        26:
        dict(name='face-2', id=25, color=[255, 255, 255], swap='face-14'),
        27:
        dict(name='face-3', id=26, color=[255, 255, 255], swap='face-13'),
        28:
        dict(name='face-4', id=27, color=[255, 255, 255], swap='face-12'),
        29:
        dict(name='face-5', id=28, color=[255, 255, 255], swap='face-11'),
        30:
        dict(name='face-6', id=29, color=[255, 255, 255], swap='face-10'),
        31:
        dict(name='face-7', id=30, color=[255, 255, 255], swap='face-9'),
        32:
        dict(name='face-8', id=31, color=[255, 255, 255], swap=''),
        33:
        dict(name='face-9', id=32, color=[255, 255, 255], swap='face-7'),
        34:
        dict(name='face-10', id=33, color=[255, 255, 255], swap='face-6'),
        35:
        dict(name='face-11', id=34, color=[255, 255, 255], swap='face-5'),
        36:
        dict(name='face-12', id=35, color=[255, 255, 255], swap='face-4'),
        37:
        dict(name='face-13', id=36, color=[255, 255, 255], swap='face-3'),
        38:
        dict(name='face-14', id=37, color=[255, 255, 255], swap='face-2'),
        39:
        dict(name='face-15', id=38, color=[255, 255, 255], swap='face-1'),
        40:
        dict(name='face-16', id=39, color=[255, 255, 255], swap='face-0'),
        41:
        dict(name='face-17', id=40, color=[255, 255, 255], swap='face-26'),
        42:
        dict(name='face-18', id=41, color=[255, 255, 255], swap='face-25'),
        43:
        dict(name='face-19', id=42, color=[255, 255, 255], swap='face-24'),
        44:
        dict(name='face-20', id=43, color=[255, 255, 255], swap='face-23'),
        45:
        dict(name='face-21', id=44, color=[255, 255, 255], swap='face-22'),
        46:
        dict(name='face-22', id=45, color=[255, 255, 255], swap='face-21'),
        47:
        dict(name='face-23', id=46, color=[255, 255, 255], swap='face-20'),
        48:
        dict(name='face-24', id=47, color=[255, 255, 255], swap='face-19'),
        49:
        dict(name='face-25', id=48, color=[255, 255, 255], swap='face-18'),
        50:
        dict(name='face-26', id=49, color=[255, 255, 255], swap='face-17'),
        51:
        dict(name='face-27', id=50, color=[255, 255, 255], swap=''),
        52:
        dict(name='face-28', id=51, color=[255, 255, 255], swap=''),
        53:
        dict(name='face-29', id=52, color=[255, 255, 255], swap=''),
        54:
        dict(name='face-30', id=53, color=[255, 255, 255], swap=''),
        55:
        dict(name='face-31', id=54, color=[255, 255, 255], swap='face-35'),
        56:
        dict(name='face-32', id=55, color=[255, 255, 255], swap='face-34'),
        57:
        dict(name='face-33', id=56, color=[255, 255, 255], swap=''),
        58:
        dict(name='face-34', id=57, color=[255, 255, 255], swap='face-32'),
        59:
        dict(name='face-35', id=58, color=[255, 255, 255], swap='face-31'),
        60:
        dict(name='face-36', id=59, color=[255, 255, 255], swap='face-45'),
        61:
        dict(name='face-37', id=60, color=[255, 255, 255], swap='face-44'),
        62:
        dict(name='face-38', id=61, color=[255, 255, 255], swap='face-43'),
        63:
        dict(name='face-39', id=62, color=[255, 255, 255], swap='face-42'),
        64:
        dict(name='face-40', id=63, color=[255, 255, 255], swap='face-47'),
        65:
        dict(name='face-41', id=64, color=[255, 255, 255], swap='face-46'),
        66:
        dict(name='face-42', id=65, color=[255, 255, 255], swap='face-39'),
        67:
        dict(name='face-43', id=66, color=[255, 255, 255], swap='face-38'),
        68:
        dict(name='face-44', id=67, color=[255, 255, 255], swap='face-37'),
        69:
        dict(name='face-45', id=68, color=[255, 255, 255], swap='face-36'),
        70:
        dict(name='face-46', id=69, color=[255, 255, 255], swap='face-41'),
        71:
        dict(name='face-47', id=70, color=[255, 255, 255], swap='face-40'),
        72:
        dict(name='face-48', id=71, color=[255, 255, 255], swap='face-54'),
        73:
        dict(name='face-49', id=72, color=[255, 255, 255], swap='face-53'),
        74:
        dict(name='face-50', id=73, color=[255, 255, 255], swap='face-52'),
        75:
        dict(name='face-51', id=74, color=[255, 255, 255], swap=''),
        76:
        dict(name='face-52', id=75, color=[255, 255, 255], swap='face-50'),
        77:
        dict(name='face-53', id=76, color=[255, 255, 255], swap='face-49'),
        78:
        dict(name='face-54', id=77, color=[255, 255, 255], swap='face-48'),
        79:
        dict(name='face-55', id=78, color=[255, 255, 255], swap='face-59'),
        80:
        dict(name='face-56', id=79, color=[255, 255, 255], swap='face-58'),
        81:
        dict(name='face-57', id=80, color=[255, 255, 255], swap=''),
        82:
        dict(name='face-58', id=81, color=[255, 255, 255], swap='face-56'),
        83:
        dict(name='face-59', id=82, color=[255, 255, 255], swap='face-55'),
        84:
        dict(name='face-60', id=83, color=[255, 255, 255], swap='face-64'),
        85:
        dict(name='face-61', id=84, color=[255, 255, 255], swap='face-63'),
        86:
        dict(name='face-62', id=85, color=[255, 255, 255], swap=''),
        87:
        dict(name='face-63', id=86, color=[255, 255, 255], swap='face-61'),
        88:
        dict(name='face-64', id=87, color=[255, 255, 255], swap='face-60'),
        89:
        dict(name='face-65', id=88, color=[255, 255, 255], swap='face-67'),
        90:
        dict(name='face-66', id=89, color=[255, 255, 255], swap=''),
        91:
        dict(name='face-67', id=90, color=[255, 255, 255], swap='face-65'),
        92:
        dict(name='left_hand_root',
             id=92,
             color=[0, 0, 255],
             swap='right_hand_root'),
        93:
        dict(name='left_thumb1', id=93, color=[0, 0, 255],
             swap='right_thumb1'),
        94:
        dict(name='left_thumb2', id=94, color=[0, 0, 255],
             swap='right_thumb2'),
        95:
        dict(name='left_thumb3', id=95, color=[0, 0, 255],
             swap='right_thumb3'),
        96:
        dict(name='left_thumb4', id=96, color=[0, 0, 255],
             swap='right_thumb4'),
        97:
        dict(name='left_forefinger1',
             id=97,
             color=[0, 0, 255],
             swap='right_forefinger1'),
        98:
        dict(name='left_forefinger2',
             id=98,
             color=[0, 0, 255],
             swap='right_forefinger2'),
        99:
        dict(name='left_forefinger3',
             id=99,
             color=[0, 0, 255],
             swap='right_forefinger3'),
        100:
        dict(name='left_forefinger4',
             id=100,
             color=[0, 0, 255],
             swap='right_forefinger4'),
        101:
        dict(name='left_middle_finger1',
             id=101,
             color=[0, 0, 255],
             swap='right_middle_finger1'),
        102:
        dict(name='left_middle_finger2',
             id=102,
             color=[0, 0, 255],
             swap='right_middle_finger2'),
        103:
        dict(name='left_middle_finger3',
             id=103,
             color=[0, 0, 255],
             swap='right_middle_finger3'),
        104:
        dict(name='left_middle_finger4',
             id=104,
             color=[0, 0, 255],
             swap='right_middle_finger4'),
        105:
        dict(name='left_ring_finger1',
             id=105,
             color=[0, 0, 255],
             swap='right_ring_finger1'),
        106:
        dict(name='left_ring_finger2',
             id=106,
             color=[0, 0, 255],
             swap='right_ring_finger2'),
        107:
        dict(name='left_ring_finger3',
             id=107,
             color=[0, 0, 255],
             swap='right_ring_finger3'),
        108:
        dict(name='left_ring_finger4',
             id=108,
             color=[0, 0, 255],
             swap='right_ring_finger4'),
        109:
        dict(name='left_pinky_finger1',
             id=109,
             color=[0, 0, 255],
             swap='right_pinky_finger1'),
        110:
        dict(name='left_pinky_finger2',
             id=110,
             color=[0, 0, 255],
             swap='right_pinky_finger2'),
        111:
        dict(name='left_pinky_finger3',
             id=111,
             color=[0, 0, 255],
             swap='right_pinky_finger3'),
        112:
        dict(name='left_pinky_finger4',
             id=112,
             color=[0, 0, 255],
             swap='right_pinky_finger4'),
        113:
        dict(name='right_hand_root',
             id=113,
             color=[0, 0, 255],
             swap='left_hand_root'),
        114:
        dict(name='right_thumb1',
             id=114,
             color=[0, 0, 255],
             swap='left_thumb1'),
        115:
        dict(name='right_thumb2',
             id=115,
             color=[0, 0, 255],
             swap='left_thumb2'),
        116:
        dict(name='right_thumb3',
             id=116,
             color=[0, 0, 255],
             swap='left_thumb3'),
        117:
        dict(name='right_thumb4',
             id=117,
             color=[0, 0, 255],
             swap='left_thumb4'),
        118:
        dict(name='right_forefinger1',
             id=118,
             color=[0, 0, 255],
             swap='left_forefinger1'),
        119:
        dict(name='right_forefinger2',
             id=119,
             color=[0, 0, 255],
             swap='left_forefinger2'),
        120:
        dict(name='right_forefinger3',
             id=120,
             color=[0, 0, 255],
             swap='left_forefinger3'),
        121:
        dict(name='right_forefinger4',
             id=121,
             color=[0, 0, 255],
             swap='left_forefinger4'),
        122:
        dict(name='right_middle_finger1',
             id=122,
             color=[0, 0, 255],
             swap='left_middle_finger1'),
        123:
        dict(name='right_middle_finger2',
             id=123,
             color=[0, 0, 255],
             swap='left_middle_finger2'),
        124:
        dict(name='right_middle_finger3',
             id=124,
             color=[0, 0, 255],
             swap='left_middle_finger3'),
        125:
        dict(name='right_middle_finger4',
             id=125,
             color=[0, 0, 255],
             swap='left_middle_finger4'),
        126:
        dict(name='right_ring_finger1',
             id=126,
             color=[0, 0, 255],
             swap='left_ring_finger1'),
        127:
        dict(name='right_ring_finger2',
             id=127,
             color=[0, 0, 255],
             swap='left_ring_finger2'),
        128:
        dict(name='right_ring_finger3',
             id=128,
             color=[0, 0, 255],
             swap='left_ring_finger3'),
        129:
        dict(name='right_ring_finger4',
             id=129,
             color=[0, 0, 255],
             swap='left_ring_finger4'),
        130:
        dict(name='right_pinky_finger1',
             id=130,
             color=[0, 0, 255],
             swap='left_pinky_finger1'),
        131:
        dict(name='right_pinky_finger2',
             id=131,
             color=[0, 0, 255],
             swap='left_pinky_finger2'),
        132:
        dict(name='right_pinky_finger3',
             id=132,
             color=[0, 0, 255],
             swap='left_pinky_finger3'),
        133:
        dict(name='right_pinky_finger4',
             id=133,
             color=[0, 0, 255],
             swap='left_pinky_finger4')
    },
    skeleton_info={
        0:
        dict(link=('neck', 'right_shoulder'), id=0, color=[255, 0, 0]),
        1:
        dict(link=('neck', 'left_shoulder'), id=1, color=[255, 85, 0]),
        2:
        dict(link=('right_shoulder', 'right_elbow'), id=2, color=[255, 170,
                                                                  0]),
        3:
        dict(link=('right_elbow', 'right_wrist'), id=3, color=[255, 255, 0]),
        4:
        dict(link=('left_shoulder', 'left_elbow'), id=4, color=[170, 255, 0]),
        5:
        dict(link=('left_elbow', 'left_wrist'), id=5, color=[85, 255, 0]),
        6:
        dict(link=('neck', 'right_hip'), id=6, color=[0, 255, 0]),
        7:
        dict(link=('right_hip', 'right_knee'), id=7, color=[0, 255, 85]),
        8:
        dict(link=('right_knee', 'right_ankle'), id=8, color=[0, 255, 170]),
        9:
        dict(link=('neck', 'left_hip'), id=9, color=[0, 255, 225]),
        10:
        dict(link=('left_hip', 'left_knee'), id=10, color=[0, 170, 255]),
        11:
        dict(link=('left_knee', 'left_ankle'), id=11, color=[0, 85, 255]),
        12:
        dict(link=('neck', 'nose'), id=12, color=[0, 0, 255]),
        13:
        dict(link=('nose', 'right_eye'), id=13, color=[255, 0, 170]),
        14:
        dict(link=('right_eye', 'right_ear'), id=14, color=[170, 0, 255]),
        15:
        dict(link=('nose', 'left_eye'), id=15, color=[255, 0, 255]),
        16:
        dict(link=('left_eye', 'left_ear'), id=16, color=[255, 0, 170]),
        17:
        dict(link=('left_hand_root', 'left_thumb1'), id=17, color=[255, 0, 0]),
        18:
        dict(link=('left_thumb1', 'left_thumb2'), id=18, color=[255, 76, 0]),
        19:
        dict(link=('left_thumb2', 'left_thumb3'), id=19, color=[255, 153, 0]),
        20:
        dict(link=('left_thumb3', 'left_thumb4'), id=20, color=[255, 230, 0]),
        21:
        dict(link=('left_hand_root', 'left_forefinger1'),
             id=21,
             color=[204, 255, 0]),
        22:
        dict(link=('left_forefinger1', 'left_forefinger2'),
             id=22,
             color=[128, 255, 0]),
        23:
        dict(link=('left_forefinger2', 'left_forefinger3'),
             id=23,
             color=[51, 255, 0]),
        24:
        dict(link=('left_forefinger3', 'left_forefinger4'),
             id=24,
             color=[0, 255, 26]),
        25:
        dict(link=('left_hand_root', 'left_middle_finger1'),
             id=25,
             color=[0, 255, 102]),
        26:
        dict(link=('left_middle_finger1', 'left_middle_finger2'),
             id=26,
             color=[0, 255, 178]),
        27:
        dict(link=('left_middle_finger2', 'left_middle_finger3'),
             id=27,
             color=[0, 255, 255]),
        28:
        dict(link=('left_middle_finger3', 'left_middle_finger4'),
             id=28,
             color=[0, 178, 255]),
        29:
        dict(link=('left_hand_root', 'left_ring_finger1'),
             id=29,
             color=[0, 102, 255]),
        30:
        dict(link=('left_ring_finger1', 'left_ring_finger2'),
             id=30,
             color=[0, 26, 255]),
        31:
        dict(link=('left_ring_finger2', 'left_ring_finger3'),
             id=31,
             color=[51, 0, 255]),
        32:
        dict(link=('left_ring_finger3', 'left_ring_finger4'),
             id=32,
             color=[128, 0, 255]),
        33:
        dict(link=('left_hand_root', 'left_pinky_finger1'),
             id=33,
             color=[204, 0, 255]),
        34:
        dict(link=('left_pinky_finger1', 'left_pinky_finger2'),
             id=34,
             color=[255, 0, 230]),
        35:
        dict(link=('left_pinky_finger2', 'left_pinky_finger3'),
             id=35,
             color=[255, 0, 153]),
        36:
        dict(link=('left_pinky_finger3', 'left_pinky_finger4'),
             id=36,
             color=[255, 0, 76]),
        37:
        dict(link=('right_hand_root', 'right_thumb1'),
             id=37,
             color=[255, 0, 0]),
        38:
        dict(link=('right_thumb1', 'right_thumb2'), id=38, color=[255, 76, 0]),
        39:
        dict(link=('right_thumb2', 'right_thumb3'), id=39, color=[255, 153,
                                                                  0]),
        40:
        dict(link=('right_thumb3', 'right_thumb4'), id=40, color=[255, 230,
                                                                  0]),
        41:
        dict(link=('right_hand_root', 'right_forefinger1'),
             id=41,
             color=[204, 255, 0]),
        42:
        dict(link=('right_forefinger1', 'right_forefinger2'),
             id=42,
             color=[128, 255, 0]),
        43:
        dict(link=('right_forefinger2', 'right_forefinger3'),
             id=43,
             color=[51, 255, 0]),
        44:
        dict(link=('right_forefinger3', 'right_forefinger4'),
             id=44,
             color=[0, 255, 26]),
        45:
        dict(link=('right_hand_root', 'right_middle_finger1'),
             id=45,
             color=[0, 255, 102]),
        46:
        dict(link=('right_middle_finger1', 'right_middle_finger2'),
             id=46,
             color=[0, 255, 178]),
        47:
        dict(link=('right_middle_finger2', 'right_middle_finger3'),
             id=47,
             color=[255, 255, 255]),
        48:
        dict(link=('right_middle_finger3', 'right_middle_finger4'),
             id=48,
             color=[0, 178, 255]),
        49:
        dict(link=('right_hand_root', 'right_ring_finger1'),
             id=49,
             color=[0, 102, 255]),
        50:
        dict(link=('right_ring_finger1', 'right_ring_finger2'),
             id=50,
             color=[0, 26, 255]),
        51:
        dict(link=('right_ring_finger2', 'right_ring_finger3'),
             id=51,
             color=[51, 0, 255]),
        52:
        dict(link=('right_ring_finger3', 'right_ring_finger4'),
             id=52,
             color=[128, 0, 255]),
        53:
        dict(link=('right_hand_root', 'right_pinky_finger1'),
             id=53,
             color=[204, 0, 255]),
        54:
        dict(link=('right_pinky_finger1', 'right_pinky_finger2'),
             id=54,
             color=[255, 0, 230]),
        55:
        dict(link=('right_pinky_finger2', 'right_pinky_finger3'),
             id=55,
             color=[255, 0, 153]),
        56:
        dict(link=('right_pinky_finger3', 'right_pinky_finger4'),
             id=56,
             color=[255, 0, 76])
    },
)

rtmlib/visualization/skeleton/openpose134.py
Tau-J/rtmlib
openpose18 = dict(
    name='openpose18',
    keypoint_info={
        0:
        dict(name='nose', id=0, color=[255, 0, 0], swap=''),
        1:
        dict(name='neck', id=1, color=[255, 85, 0], swap=''),
        2:
        dict(name='right_shoulder',
             id=2,
             color=[255, 170, 0],
             swap='left_shoulder'),
        3:
        dict(name='right_elbow', id=3, color=[255, 255, 0], swap='left_elbow'),
        4:
        dict(name='right_wrist', id=4, color=[170, 255, 0], swap='left_wrist'),
        5:
        dict(name='left_shoulder',
             id=5,
             color=[85, 255, 0],
             swap='right_shoulder'),
        6:
        dict(name='left_elbow', id=6, color=[0, 255, 0], swap='right_elbow'),
        7:
        dict(name='left_wrist', id=7, color=[0, 255, 85], swap='right_wrist'),
        8:
        dict(name='right_hip', id=8, color=[0, 255, 170], swap='left_hip'),
        9:
        dict(name='right_knee', id=9, color=[0, 255, 255], swap='left_knee'),
        10:
        dict(name='right_ankle', id=10, color=[0, 170, 255],
             swap='left_ankle'),
        11:
        dict(name='left_hip', id=11, color=[0, 85, 255], swap='right_hip'),
        12:
        dict(name='left_knee', id=12, color=[0, 0, 255], swap='right_knee'),
        13:
        dict(name='left_ankle', id=13, color=[85, 0, 255], swap='right_ankle'),
        14:
        dict(name='right_eye', id=14, color=[170, 0, 255], swap='left_eye'),
        15:
        dict(name='left_eye', id=15, color=[255, 0, 255], swap='right_eye'),
        16:
        dict(name='right_ear', id=16, color=[255, 0, 170], swap='left_ear'),
        17:
        dict(name='left_ear', id=17, color=[255, 0, 85], swap='right_ear'),
    },
    skeleton_info={
        0: dict(link=('neck', 'right_shoulder'), id=0, color=[255, 0, 0]),
        1: dict(link=('neck', 'left_shoulder'), id=1, color=[255, 85, 0]),
        2: dict(link=('right_shoulder', 'right_elbow'),
                id=2,
                color=[255, 170, 0]),
        3: dict(link=('right_elbow', 'right_wrist'), id=3, color=[255, 255,
                                                                  0]),
        4: dict(link=('left_shoulder', 'left_elbow'),
                id=4,
                color=[170, 255, 0]),
        5: dict(link=('left_elbow', 'left_wrist'), id=5, color=[85, 255, 0]),
        6: dict(link=('neck', 'right_hip'), id=6, color=[0, 255, 0]),
        7: dict(link=('right_hip', 'right_knee'), id=7, color=[0, 255, 85]),
        8: dict(link=('right_knee', 'right_ankle'), id=8, color=[0, 255, 170]),
        9: dict(link=('neck', 'left_hip'), id=9, color=[0, 255, 225]),
        10: dict(link=('left_hip', 'left_knee'), id=10, color=[0, 170, 255]),
        11: dict(link=('left_knee', 'left_ankle'), id=11, color=[0, 85, 255]),
        12: dict(link=('neck', 'nose'), id=12, color=[0, 0, 255]),
        13: dict(link=('nose', 'right_eye'), id=13, color=[255, 0, 170]),
        14: dict(link=('right_eye', 'right_ear'), id=14, color=[170, 0, 255]),
        15: dict(link=('nose', 'left_eye'), id=15, color=[255, 0, 255]),
        16: dict(link=('left_eye', 'left_ear'), id=16, color=[255, 0, 170]),
    },
)

rtmlib/visualization/skeleton/openpose18.py
Tau-J/rtmlib
from setuptools import find_packages, setup


def readme():
    with open('README.md', encoding='utf-8') as f:
        content = f.read()
    return content


version_file = 'rtmlib/version.py'


def get_version():
    with open(version_file, 'r') as f:
        exec(compile(f.read(), version_file, 'exec'))
    import sys

    # return short version for sdist
    if 'sdist' in sys.argv or 'bdist_wheel' in sys.argv:
        return locals()['short_version']
    else:
        return locals()['__version__']


def parse_requirements(fname='requirements.txt', with_version=True):
    """Parse the package dependencies listed in a requirements file but strips
    specific versioning information.

    Args:
        fname (str): path to requirements file
        with_version (bool, default=False): if True include version specs

    Returns:
        List[str]: list of requirements items

    CommandLine:
        python -c "import setup; print(setup.parse_requirements())"
    """
    import re
    import sys
    from os.path import exists
    require_fpath = fname

    def parse_line(line):
        """Parse information from a line in a requirements text file."""
        if line.startswith('-r '):
            # Allow specifying requirements in other files
            target = line.split(' ')[1]
            for info in parse_require_file(target):
                yield info
        else:
            info = {'line': line}
            if line.startswith('-e '):
                info['package'] = line.split('#egg=')[1]
            elif '@git+' in line:
                info['package'] = line
            else:
                # Remove versioning from the package
                pat = '(' + '|'.join(['>=', '==', '>']) + ')'
                parts = re.split(pat, line, maxsplit=1)
                parts = [p.strip() for p in parts]

                info['package'] = parts[0]
                if len(parts) > 1:
                    op, rest = parts[1:]
                    if ';' in rest:
                        # Handle platform specific dependencies
                        # http://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-platform-specific-dependencies
                        version, platform_deps = map(str.strip,
                                                     rest.split(';'))
                        info['platform_deps'] = platform_deps
                    else:
                        version = rest  # NOQA
                    info['version'] = (op, version)
            yield info

    def parse_require_file(fpath):
        with open(fpath, 'r') as f:
            for line in f.readlines():
                line = line.strip()
                if line and not line.startswith('#'):
                    for info in parse_line(line):
                        yield info

    def gen_packages_items():
        if exists(require_fpath):
            for info in parse_require_file(require_fpath):
                parts = [info['package']]
                if with_version and 'version' in info:
                    parts.extend(info['version'])
                if not sys.version.startswith('3.4'):
                    # apparently package_deps are broken in 3.4
                    platform_deps = info.get('platform_deps')
                    if platform_deps is not None:
                        parts.append(';' + platform_deps)
                item = ''.join(parts)
                yield item

    packages = list(gen_packages_items())
    return packages


if __name__ == '__main__':
    setup(name='rtmlib',
          version=get_version(),
          description='A library for real-time pose estimation.',
          author='Tau-J',
          author_email='taujiang@outlook.com',
          keywords='pose estimation',
          long_description=readme(),
          long_description_content_type='text/markdown',
          packages=find_packages(),
          include_package_data=True,
          url='https://github.com/Tau-J/rtmlib',
          license='Apache License 2.0',
          python_requires='>=3.7',
          install_requires=parse_requirements('requirements/runtime.txt'),
          extras_require={
              'all': parse_requirements('requirements.txt'),
              'optional': parse_requirements('requirements/optional.txt'),
          })

setup.py
Tau-J/rtmlib
import gradio as gr
import numpy as np

from rtmlib import Body, Wholebody, draw_skeleton

cached_model = {}


def predict(img,
            openpose_skeleton,
            model_type,
            black_bg=False,
            backend='onnxruntime',
            device='cpu'):

    if model_type == 'body':
        constructor = Body
    elif model_type == 'wholebody':
        constructor = Wholebody
    else:
        raise NotImplementedError

    model_id = str((constructor.__qualname__, openpose_skeleton, black_bg,
                    backend, device))

    if model_id in cached_model:
        model = cached_model[model_id]
    else:
        model = constructor(to_openpose=openpose_skeleton,
                            backend=backend,
                            device=device)
        cached_model[model_id] = model

    keypoints, scores = model(img)

    if black_bg:
        img_show = np.zeros_like(img, dtype=np.uint8)
    else:
        img_show = img.copy()

    img_show = draw_skeleton(img_show,
                             keypoints,
                             scores,
                             openpose_skeleton=openpose_skeleton,
                             kpt_thr=0.4)
    return img_show[:, :, ::-1]


with gr.Blocks() as demo:

    with gr.Tab('Upload-Image'):
        input_img = gr.Image(type='numpy')
        button = gr.Button('Inference', variant='primary')
        openpose_skeleton = gr.Checkbox(label='openpose-skeleton',
                                        info='Draw OpenPose-style Skeleton')
        black_bg = gr.Checkbox(
            label='black background',
            info='Whether to draw skeleton on black background')
        model_type = gr.Dropdown(['body', 'wholebody'],
                                 label='Keypoint Type',
                                 info='Body / Wholebody',
                                 value='body')
        backend = gr.Dropdown(['opencv', 'onnxruntime', 'openvino'],
                              label='Choose backend',
                              info='opencv / onnxruntime / openvino',
                              value='onnxruntime')
        device = gr.Dropdown(['cpu', 'cuda'],
                             label='Choose device',
                             info='cpu / cuda',
                             value='cpu')

        gr.Markdown('## Output')
        out_image = gr.Image(type='numpy')
        gr.Examples(['./demo.jpg'], input_img)

        button.click(predict, [
            input_img, openpose_skeleton, model_type, black_bg, backend, device
        ], out_image)

gr.close_all()
demo.queue()
demo.launch()

webui.py
Tau-J/rtmlib
import time

import cv2

from rtmlib import PoseTracker, Wholebody, draw_skeleton

# import numpy as np

device = 'cpu'
backend = 'onnxruntime'  # opencv, onnxruntime, openvino

cap = cv2.VideoCapture('./demo.mp4')

openpose_skeleton = False  # True for openpose-style, False for mmpose-style

wholebody = PoseTracker(
    Wholebody,
    det_frequency=10,
    to_openpose=openpose_skeleton,
    mode='balanced',  # balanced, performance, lightweight
    backend=backend,
    device=device)

frame_idx = 0

while cap.isOpened():
    success, frame = cap.read()
    frame_idx += 1

    if not success:
        break
    s = time.time()
    keypoints, scores = wholebody(frame)
    det_time = time.time() - s
    print('det: ', det_time)

    img_show = frame.copy()

    # if you want to use black background instead of original image,
    # img_show = np.zeros(img_show.shape, dtype=np.uint8)

    img_show = draw_skeleton(img_show,
                             keypoints,
                             scores,
                             openpose_skeleton=openpose_skeleton,
                             kpt_thr=0.43)

    img_show = cv2.resize(img_show, (960, 540))
    cv2.imshow('img', img_show)
    cv2.waitKey(10)

wholebody_demo.py
Tau-J/rtmlib