update

README.md

@@ -7,6 +7,7 @@ Basically, rtmlib only requires these dependencies:
 - numpy
 - opencv-python
 - opencv-contrib-python
+- onnxruntime
 
 Optionally, you can use other common backends like opencv, onnxruntime, tensorrt to accelerate the inference process.
 
@@ -29,8 +30,6 @@ pip install -r requirements.txt
 pip install -e .
 
 # [optional]
-# pip install onnxruntime
-# or
 # pip install onnxruntime-gpu
 ```
 
@@ -54,7 +53,7 @@ import cv2
 
 from rtmlib import Wholebody, draw_skeleton
 
-device = 'cpu'
+device = 'cpu'  # cpu, cuda
 backend = 'onnxruntime'  # opencv, onnxruntime
 img = cv2.imread('./demo.jpg')
 

requirements.txt

@@ -1,3 +1,4 @@
 numpy
+onnxruntime
 opencv-contrib-python
 opencv-python

rtmlib/tools/base.py

@@ -37,9 +37,17 @@ def __init__(self,
         providers = RTMLIB_SETTINGS[backend][device]
 
         if backend == 'opencv':
-            session = cv2.dnn.readNetFromONNX(onnx_model)
-            session.setPreferableBackend(providers[0])
-            session.setPreferableTarget(providers[1])
+            try:
+                session = cv2.dnn.readNetFromONNX(onnx_model)
+                session.setPreferableBackend(providers[0])
+                session.setPreferableTarget(providers[1])
+            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

rtmlib/tools/object_detection/rtmdet.py

@@ -12,7 +12,7 @@ class RTMDet(BaseTool):
     def __init__(self,
                  onnx_model: str = 'rtmdet-m-640x640',
                  model_input_size: tuple = (640, 640),
-                 backend: str = 'opencv',
+                 backend: str = 'onnxruntime',
                  device: str = 'cpu'):
         super().__init__(onnx_model, model_input_size, backend, device)
         raise NotImplementedError
@@ -105,7 +105,7 @@ def postprocess(
 
         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

rtmlib/tools/object_detection/yolox.py

@@ -15,7 +15,7 @@ def __init__(self,
                  model_input_size: tuple = (640, 640),
                  nms_thr=0.45,
                  score_thr=0.7,
-                 backend: str = 'opencv',
+                 backend: str = 'onnxruntime',
                  device: str = 'cpu'):
         super().__init__(onnx_model,
                          model_input_size,
@@ -112,7 +112,7 @@ def postprocess(
 
         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

rtmlib/tools/pose_estimation/rtmpose.py

@@ -15,7 +15,7 @@ def __init__(self,
                  mean: tuple = (123.675, 116.28, 103.53),
                  std: tuple = (58.395, 57.12, 57.375),
                  to_openpose: bool = False,
-                 backend: str = 'opencv',
+                 backend: str = 'onnxruntime',
                  device: str = 'cpu'):
         super().__init__(onnx_model, model_input_size, mean, std, backend,
                          device)

rtmlib/tools/solution/body.py

@@ -38,7 +38,7 @@ def __init__(self,
                  pose_input_size: tuple = (288, 384),
                  mode: str = 'performance',
                  to_openpose: bool = False,
-                 backend: str = 'opencv',
+                 backend: str = 'onnxruntime',
                  device: str = 'cpu'):
 
         if det is None:

rtmlib/tools/solution/pose_tracker.py

@@ -0,0 +1,67 @@
+import numpy as np
+
+from .. import YOLOX, RTMPose
+
+
+class PoseTracker:
+
+    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 = 'performance',
+                 to_openpose: bool = False,
+                 backend: str = 'opencv',
+                 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/wholebody.py

@@ -42,7 +42,7 @@ def __init__(self,
                  pose_input_size: tuple = (288, 384),
                  mode: str = 'performance',
                  to_openpose: bool = False,
-                 backend: str = 'opencv',
+                 backend: str = 'onnxruntime',
                  device: str = 'cpu'):
 
         if det is None:

rtmlib/version.py

@@ -1,4 +1,4 @@
-__version__ = '0.0.1'
+__version__ = '0.0.4'
 short_version = __version__