diff --git a/README.md b/README.md index 0ef5a31..b2eeb95 100644 --- a/README.md +++ b/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') diff --git a/requirements.txt b/requirements.txt index e89b2dd..3b07038 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ numpy +onnxruntime opencv-contrib-python opencv-python diff --git a/rtmlib/tools/base.py b/rtmlib/tools/base.py index 5a78a05..215f56b 100644 --- a/rtmlib/tools/base.py +++ b/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 diff --git a/rtmlib/tools/object_detection/rtmdet.py b/rtmlib/tools/object_detection/rtmdet.py index b11c589..b045503 100644 --- a/rtmlib/tools/object_detection/rtmdet.py +++ b/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 diff --git a/rtmlib/tools/object_detection/yolox.py b/rtmlib/tools/object_detection/yolox.py index a023129..38bc57b 100644 --- a/rtmlib/tools/object_detection/yolox.py +++ b/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 diff --git a/rtmlib/tools/pose_estimation/rtmpose.py b/rtmlib/tools/pose_estimation/rtmpose.py index 9807e38..26e11e9 100644 --- a/rtmlib/tools/pose_estimation/rtmpose.py +++ b/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) diff --git a/rtmlib/tools/solution/body.py b/rtmlib/tools/solution/body.py index 9ed77b2..9d152de 100644 --- a/rtmlib/tools/solution/body.py +++ b/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: diff --git a/rtmlib/tools/solution/pose_tracker.py b/rtmlib/tools/solution/pose_tracker.py new file mode 100644 index 0000000..734d20c --- /dev/null +++ b/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 diff --git a/rtmlib/tools/solution/wholebody.py b/rtmlib/tools/solution/wholebody.py index 5a4465d..3cdd116 100644 --- a/rtmlib/tools/solution/wholebody.py +++ b/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: diff --git a/rtmlib/version.py b/rtmlib/version.py index 3406aa5..bd7ec83 100644 --- a/rtmlib/version.py +++ b/rtmlib/version.py @@ -1,4 +1,4 @@ -__version__ = '0.0.1' +__version__ = '0.0.4' short_version = __version__