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

OpenCV学堂

发布于 2023-08-22 13:20:28

9920

发布于 2023-08-22 13:20:28

文章被收录于专栏：贾志刚-OpenCV学堂

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

前言

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

程序实现

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

直接通过PT模型推理的好处有两点，一个是不需要写部署代码了，二是精度不会在模型转化中有细微损失了。特别适合Python开发者

对象检测 - 运行截图如下：