rtmlib

Use the PoseTracker class to perform pose estimation on a continuous video stream or file. PoseTracker manages detection frequency via det_frequency (detecting every N frames) to optimize performance.

from rtmlib import Body, Custom, PoseTracker, draw_skeleton
import cv2

cap = cv2.VideoCapture(0)  # Use cap = cv2.VideoCapture('./demo.mp4') for video files

device = 'cpu'
backend = 'onnxruntime'
openpose_skeleton = False

pose_tracker = PoseTracker(Body,
                        mode='balanced',
                        det_frequency=10,  # detect every 10 frames
                        backend=backend, device=device,
                        to_openpose=False)

frame_idx = 0
while cap.isOpened():
    success, frame = cap.read()
    frame_idx += 1
    if not success:
        break

    keypoints, scores = pose_tracker(frame)

    img_show = frame.copy()
    img_show = draw_skeleton(img_show,
                             keypoints,
                             scores,
                             openpose_skeleton=openpose_skeleton,
                             kpt_thr=0.43)
    cv2.imshow('img', img_show)
    cv2.waitKey(10)

Use the Wholebody class for high-level pose estimation on a single image. You can specify the device (cpu, cuda, or mps), the backend (opencv, onnxruntime, or openvino), and the mode (performance, lightweight, or balanced). Use draw_skeleton to visualize the results. For animal pose estimation, set to_openpose=True.

import cv2
from rtmlib import Wholebody, draw_skeleton

img = cv2.imread('./demo.jpg')

device = 'cpu'  # cpu, cuda, mps
backend = 'onnxruntime'  # opencv, onnxruntime, openvino
openpose_skeleton = False  # True for openpose-style (required for animals), False for mmpose-style

wholebody = Wholebody(to_openpose=openpose_skeleton,
                      mode='balanced',  # 'performance', 'lightweight', 'balanced'. Default: 'balanced'
                      backend=backend, device=device)
keypoints, scores = wholebody(img)

# visualize
img = draw_skeleton(img, keypoints, scores, kpt_thr=0.5, to_openpose=openpose_skeleton)
cv2.imshow('img', img)
cv2.waitKey(0)

High-level APIs (Solutions) allow you to specify models using either a mode string or by providing specific det (detector) and pose (estimator) paths/URLs.

Using mode:

from rtmlib import Wholebody
wholebody = Wholebody(mode='performance', backend=backend, device=device)

Using specific det and pose paths:

from rtmlib import Body
body = Body(det='path/to/detector.zip', 
            det_input_size=(640, 640), 
            pose='path/to/pose.zip', 
            pose_input_size=(288, 384), 
            backend=backend, 
            device=device)

Using Custom for multiclass or specific configurations:

from rtmlib import Custom
custom = Custom(det_class='YOLOX', 
               det_mode='multiclass', 
               det='path/to/detector.onnx', 
               det_input_size=(640, 640), 
               pose_class='ViTPose', 
               pose='path/to/pose.onnx', 
               pose_input_size=(192, 256), 
               backend=backend, 
               device=device)

The rtmlib Model Zoo provides pre-trained ONNX models for various detection tasks. Models are categorized by the type of object they detect:

• Person: Includes models trained on COCO (real humans only) and HumanArt (real humans and cartoon characters). Examples include YOLOX-nano, YOLOX-tiny, YOLOX-s, YOLOX-m, YOLOX-l, and YOLOX-x.
• Hand: Includes RTMDet-nano trained on 5 datasets.
• Multi-class: Uses YOLOX models (nano, tiny, s, m, l, Darknet53, X) trained on COCO classes, covering 80 categories such as person, bicycle, car, dog, etc.

COCO_CLASSES = [
    'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
    'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
    'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
    'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
    'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
    'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
    'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
    'scissors', 'teddy bear', 'hair drier', 'toothbrush']