两行代码搞定YOLOv8各种模型推理

新闻资讯 2023-06-15 22:23 112 0

点击上方蓝字关注我们

微信公众号：OpenCV学堂
关注获取更多计算机视觉与深度学习知识

前言

大家好，YOLOv8 框架本身提供的API函数是可以两行代码实现 YOLOv8 模型推理，这次我把这段代码封装成了一个类，只有40行代码左右，可以同时支持YOLOv8对象检测、实例分割、姿态评估模型的GPU与CPU上推理演示。

程序实现

使用PyQT5开发一个简单的YOLOv8 框架本身提供的API函数演示交互界面，支持从界面上选择模型文件、测试图像或者视频文件，点击开始推理按钮就可以运行了，整个程序的界面如下：

YOLOv8框架支持的函数推理会自动识别模型的类型是对象检测、实例分割、姿态评估中哪一种，有GPU支持的情况下，默认会使用GPU推理。

推理运行在一个单独的PyQT线程中，通过信号与槽机制实现推理结果返回与更新。实现的线程代码如下：

class InferenceThread(QtCore.QThread):
    fire_stats_signal = QtCore.pyqtSignal(dict)

    def __init__(self, settings):
        super(InferenceThread, self).__init__()
        self.settings = settings
        self.detector = YOLOv8PtInference(settings)
        self.input_image = settings.input_image

    def run(self):
        if self.detector is None:
            return
        if self.input_image.endswith(".mp4"):
            cap = cv.VideoCapture(self.input_image)
            while True:
                ret, frame = cap.read()
                if ret is True:
                    self.detector.infer_image(frame)
                    self.fire_stats_signal.emit({"result": frame})
                else:
                    break
        else:
            frame = cv.imread(self.input_image)
            self.detector.infer_image(frame)
            self.fire_stats_signal.emit({"result": frame})
        self.fire_stats_signal.emit({"done": "done"})
        return