add naive pose tracker

README.md

@@ -1,5 +1,7 @@
 # rtmlib
 
+![demo](https://github.com/Tau-J/rtmlib/assets/13503330/b7e8ce8b-3134-43cf-bba6-d81656897289)
+
 rtmlib is a super lightweight library to conduct pose estimation based on [RTMPose](https://github.com/open-mmlab/mmpose/tree/dev-1.x/projects/rtmpose) models **WITHOUT** any dependencies like mmcv, mmpose, mmdet, etc.
 
 Basically, rtmlib only requires these dependencies:
@@ -82,6 +84,7 @@ cv2.waitKey()
 - Solutions (High-level APIs)
   - [Wholebody](/rtmlib/tools/solution/wholebody.py)
   - [Body](/rtmlib/tools/solution/body.py)
+  - [PoseTracker](/rtmlib/tools/solution/pose_tracker.py)
 - Models (Low-level APIs)
   - [YOLOX](/rtmlib/tools/object_detection/yolox.py)
   - [RTMDet](/rtmlib/tools/object_detection/rtmdet.py)

rtmlib/__init__.py

@@ -1,7 +1,7 @@
 from .draw import draw_bbox, draw_skeleton
-from .tools import YOLOX, Body, RTMDet, RTMPose, Wholebody
+from .tools import YOLOX, Body, PoseTracker, RTMDet, RTMPose, Wholebody
 
 __all__ = [
     'RTMDet', 'RTMPose', 'YOLOX', 'Wholebody', 'Body', 'draw_skeleton',
-    'draw_bbox'
+    'draw_bbox', 'PoseTracker'
 ]

rtmlib/tools/__init__.py

@@ -1,5 +1,5 @@
 from .object_detection import YOLOX, RTMDet
 from .pose_estimation import RTMPose
-from .solution import Body, Wholebody
+from .solution import Body, PoseTracker, Wholebody
 
-__all__ = ['RTMDet', 'RTMPose', 'YOLOX', 'Wholebody', 'Body']
+__all__ = ['RTMDet', 'RTMPose', 'YOLOX', 'Wholebody', 'Body', 'PoseTracker']

rtmlib/tools/solution/__init__.py

@@ -1,4 +1,5 @@
 from .body import Body
+from .pose_tracker import PoseTracker
 from .wholebody import Wholebody
 
-__all__ = ['Wholebody', 'Body']
+__all__ = ['Wholebody', 'Body', 'PoseTracker']

rtmlib/tools/solution/pose_tracker.py

@@ -1,3 +1,44 @@
+'''
+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
 
 
@@ -23,6 +64,16 @@ def pose_to_bbox(keypoints: np.ndarray, expansion: float = 1.25) -> np.ndarray:
 
 
 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,

rtmlib/tools/solution/wholebody.py

@@ -3,7 +3,7 @@
 import numpy as np
 
 from .. import YOLOX, RTMPose
-from .types import BodyResult, Keypoint, PoseResult
+from .utils.types import BodyResult, Keypoint, PoseResult
 
 
 class Wholebody:

wholebody_demo.py

@@ -6,19 +6,17 @@
 
 # import numpy as np
 
-device = 'cpu'
+device = 'cuda'
 backend = 'onnxruntime'  # opencv, onnxruntime
 
-cap = cv2.VideoCapture('./demo.jpg')
+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)
 
-video_writer = None
-pred_instances_list = []
 frame_idx = 0
 
 while cap.isOpened():
@@ -45,4 +43,4 @@
 
     img_show = cv2.resize(img_show, (960, 540))
     cv2.imshow('img', img_show)
-    cv2.waitKey()
+    cv2.waitKey(10)