openvino yolo11 pose推理代码

import openvino as ov
import cv2
import pathlib
from ultralytics import YOLO
from typing import Union


class PoseDetector:
    def __init__(
            self,
            model_path: Union[str, pathlib.Path],
            device: str = "CPU",
            conf_threshold: float = 0.5,
            input_size: tuple = (640, 640)
    ):
        self.core = ov.Core()
        self.device = device
        self.conf_threshold = conf_threshold
        self.input_size = input_size
        self.model_path = pathlib.Path(model_path)

        # 初始化模型
        self._load_models()
        self._init_predictor()

        # 视频捕获对象
        self.cap = None

    def _load_models(self):
        """加载OpenVINO模型并进行编译"""
        # 加载OpenVINO模型
        ov_model = self.core.read_model(self.model_path)

        # 设备特定配置
        ov_config = {}
        if self.device != "CPU":
            ov_model.reshape({0: [1, 3, *self.input_size]})
        if "GPU" in self.device or ("AUTO" in self.device and "GPU" in self.core.available_devices):
            ov_config["GPU_DISABLE_WINOGRAD_CONVOLUTION"] = "YES"

        self.compiled_model = self.core.compile_model(ov_model, self.device, ov_config)

    def _init_predictor(self):
        self.pose_model = YOLO(self.model_path.parent, task="pose")
        if self.pose_model.predictor is None:
            custom_args = {
                "conf": self.conf_threshold,
                "batch": 1,
                "save": False,
                "mode": "predict"
            }
            self.pose_model.predictor = self.pose_model._smart_load("predictor")(
                overrides={**self.pose_model.overrides, **custom_args},
                _callbacks=self.pose_model.callbacks
            )
            self.pose_model.predictor.setup_model(model=self.pose_model.model)

        # 绑定已编译的OpenVINO模型
        self.pose_model.predictor.model.ov_compiled_model = self.compiled_model

    def process_frame(self, frame):
        results = self.pose_model(frame)
        return results[0].plot()[:, :, ::-1]  # 转换颜色通道顺序


    def run(self, video_source: Union[str, int] = 0):
        self.cap = cv2.VideoCapture(video_source)
        while self.cap.isOpened():
            ret, frame = self.cap.read()
            if not ret:
                break
            processed_frame = self.process_frame(frame)
            cv2.cvtColor(processed_frame,cv2.COLOR_BGR2RGB)
            cv2.imshow("Pose Detection", processed_frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        self.release()

    def release(self):
        if self.cap:
            self.cap.release()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    # 使用示例
    detector = PoseDetector(
        model_path="D:/yolov11-master/weights/yolo11n-pose_openvino_model/yolo11n-pose.xml",
        device="CPU",
        conf_threshold=0.25
    )
    detector.run(0)  # 使用摄像头

模型转换的话，export.py  改成openvino就好了