【v5lite】调用onnx推理

前言

跟着博主导入的加以修改的，反正v5lite的版本要是1.4版本的，不然容易出现错误！
后面再去把博主的博文导进来
树莓派4B运行yolov5lite转onnx模型（实测3-7fps）这个比较浅显易懂，跟着博主搞几乎不会出现错误。

还有的参考，因为个人水平有限，我本人没看懂，但是我认为极具参考价值，一致贴上来！
1、pogg源大佬的知乎文章工程部署（五）：这周不鸽！v5Lite树莓派15帧教程
但是我没看懂这里QWQ转换出onnx模型之后，我不知道如何检测模型的好坏
2、基于树莓派4B的YOLOv5-Lite目标检测的移植与部署（含训练教程）

一、主线程

def main():
    # 模型加载
    model_pb_path = "D:\\YOLOv5-Lite-1.4\\runs\\train\\exp1\\weights\\best.onnx"
    so = ort.SessionOptions()
    net = ort.InferenceSession(model_pb_path, so)
    """
    这里首先指定了模型文件的路径（model_pb_path），然后创建了一个ONNX运行时的会话选项（so），
    最后使用模型路径和会话选项来初始化一个推理会话（net），用于后续对模型进行推理操作。
    """

    # 标签字典
    global dic_labels
    dic_labels = {
   0: 'good_apple',
                  1: 'bad_apple',
                  2: 'good_banana',
                  3: 'bad_banana',
                  4: 'good_mango',
                  5: 'bad_mango'}
    """
    定义了一个全局的字典（dic_labels），用于将模型输出的类别索引映射到对应的水果标签，
    方便后续在检测结果中显示有意义的水果名称。
    """

    # 模型参数
    model_h = 320
    model_w = 320
    nl = 3
    na = 3
    stride = [8., 16., 32.]
    anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]]
    anchor_grid = np.asarray(anchors, dtype=np.float32).reshape(nl, -1, 2)
    """
    这里定义了一系列模型相关的参数：
    - model_h和model_w指定了模型输入图像的高度和宽度。
    - nl表示模型的检测层数。
    - na表示每层的锚框数量。
    - stride是不同检测层的步长信息。
    - anchors是定义的锚框尺寸列表，经过处理后转换为numpy数组并重塑为合适的形状存储在anchor_grid中，
      这些参数在目标检测模型中用于确定检测框的位置和大小等。
    """

    # 初始化队列和线程
    img_queue = queue.Queue()
    result_queue = queue.Queue()
    result_copy_queue = queue.Queue()
    detection_thread = DetectionThread(img_queue, result_queue, net, model_h, model_w, nl, na, stride, anchor_grid, dic_labels)
    detection_thread.start()
    """
    创建了三个队列：
    - img_queue用于存储待检测的图像数据。
    - result_queue用于存储检测线程得到的检测结果。
    - result_copy_queue用于在另一个线程（串口线程）中传递检测结果副本。

    然后创建了一个检测线程（detection_thread），并传入相关参数后启动该线程，
    该线程将从img_queue中获取图像进行检测，并将结果放入result_queue。
    """

    # 串口线程配置
    serial_port = 'COM12'
    baud_rate = 115200
    serial_thread = SerialThread(serial_port, baud_rate, result_copy_queue)
    serial_thread.start()
    """
    配置了串口相关的参数，包括串口号（serial_port）和波特率（baud_rate），
    然后创建了一个串口线程（serial_thread），并传入相关参数后启动该线程，
    该线程将从result_copy_queue中获取检测结果副本进行后续处理（可能是通过串口发送数据等操作）。
    """

    video = 0
    cap = cv2.VideoCapture(video)
    flag_det = False
    """
    设置视频源为默认摄像头（video = 0），然后使用cv2.VideoCapture创建一个视频捕获对象（cap），
    用于读取视频帧。同时初始化了一个标志变量flag_det为False，用于控制是否启动检测功能。
    """

    prev_time = time.time()  # 上一帧的时间戳
    frame_count = 0          # 总帧数
    try:
        while True:
            success, img0 = cap.read()
            if not success:
                break
            curr_time = time.time()
            # 将图像放入队列（即使检测未启动）
            img_queue.put(img0)
            """
            在循环中，不断从视频捕获对象中读取视频帧（img0），如果读取成功，
            将当前帧放入img_queue队列中，无论检测功能是否启动，都先将图像放入队列等待后续处理。
            """

            # 如果检测启动，从结果队列中获取结果并绘制
            if flag_det:
                if not result_queue.empty():
                    det_boxes, scores, ids = result_queue.get()
                    for box, score, id in zip(det_boxes, scores, ids):
                        label = '%s:%.2f' % (dic_labels[id], score)
                        plot_one_box(box.astype(np.int16), img0, color=(255, 0, 0), label=label, line_thickness=None)
                    result_copy_queue.put( (det_boxes, scores, ids))        # 复制线程
            """
            如果检测标志flag_det为True，表示检测功能已启动，当result_queue队列中有检测结果时，
            从队列中取出检测框信息（det_boxes）、得分（scores）和类别索引（ids），
            然后遍历这些结果，根据类别索引从dic_labels字典中获取对应的水果标签，
            并在图像上绘制检测框和标注信息，最后将检测结果副本放入result_copy_queue队列，以便串口线程获取。
            """

            frame_count += 1
            if frame_count > 1:  # 避免除以零错误
                fps = 1. / (curr_time - prev_time)
                str_FPS = "FPS: %.2f" % fps
            else:
                str_FPS = "FPS: N/A"  # 第一帧时显示N/A
            prev_time = curr_time
            cv2.putText(img0, str_FPS, (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 3)
            """
            每处理一帧图像，帧数（frame_count）加1，根据当前帧和上一帧的时间戳计算帧率（fps），
            并将帧率信息以文本形式绘制在图像上，方便实时查看处理速度。
            """

            cv2.imshow("FruitSeeSee", img0)
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break
            elif key == ord('s'):
                flag_det = not flag_det
                if flag_det:
                    print("Detection started.")
                else:
                    print("Detection stopped.")
                    # 清空队列以防有残留数据
                    while not result_queue.empty():
                        result_queue.get()
                    while not img_queue.empty():
                        img_queue.get()
            """
            显示处理后的图像（img0），并等待用户按键操作：
            - 如果按下'q'键，退出循环，结束程序。
            - 如果按下's'键，切换检测功能的启动和停止状态，当启动检测时打印相应提示信息，
              当停止检测时，除了打印提示信息，还会清空result_queue和img_queue队列中的残留数据。
            """

    finally:
        detection_thread.stop()
        detection_thread.join()
        serial_thread.stop()
        serial_thread.join()
        cap.release()
        cv2.destroyAllWindows()
        """
        在程序结束时（无论是否正常结束），执行以下清理操作：
        - 停止并等待检测线程（detection_thread）结束。
        - 停止并等待串口线程（serial_thread）结束。
        - 释放视频捕获对象（cap）占用的资源。
        - 关闭所有打开的cv2窗口。
        """

二、推理线程

识别线程

"""################# 识别线程 ###########################"""
class DetectionThread(threading.Thread):
    def __init__(self, img_queue, result_queue, net, model_h, model_w, nl, na, stride, anchor_grid, dic_labels):
        """
        识别线程类的构造函数，用于初始化线程相关的属性。

        :param img_queue: 存储待检测图像的队列
        :param result_queue: 用于存储检测结果的队列
        :param net: ONNX模型推理会话对象，用于对图像进行推理
        :param model_h: 模型输入图像的高度
        :param model_w: 模型输入图像的宽度
        :param nl: 模型的检测层数
        :param na: 每层的锚框数量
        :param stride: 不同检测层的步长信息
        :param anchor_grid: 处理后的锚框尺寸数组，用于辅助检测
        :param dic_labels: 类别标签字典，将类别索引映射到具体的标签名称
        """
        threading.Thread.__init__(self)
        self.img_queue = img_queue
        self.result_queue = result_queue
        self.net = net
        self.model_h = model_h
        self.model_w = model_w
        self.nl = nl
        self.na = na
        self.stride = stride
        self.anchor_grid = anchor_grid
        self.dic_labels = dic_labels
        self.running = False

    def run(self):
        """
        线程启动后执行的方法，在此方法中实现了图像的检测和结果推送逻辑。

        只要线程处于运行状态（self.running为True），就会不断从图像队列中获取图像进行检测，
        并将检测结果放入结果队列。
        """
        self.running = True
        while self.rrunning:
            if not self.img_queue.empty():
                img0 = self.img_queue.get()
                det_boxes, scores, ids = infer_img(img0, self.net, self.model_h, self.model_w, self.nl, self.na, self.stride, self.anchor_grid)
                self.result_queue.put((det_boxes, scores, ids))
            time.sleep(0.01)  # 防止线程过于繁忙

    def stop(self):
        """
        用于停止线程运行的方法，通过将self.running设置为False，
        使得线程在下次循环检查时退出循环，从而停止线程的执行。
        """
        self.running = False

重头戏
plot_one_box函数：
功能：在给定图像上绘制一个检测框及对应的标签。
实现：首先根据图像大小动态计算绘制检测框的线条粗细（tl），如果未指定颜色则随机生成一个颜色。然后根据传入的检测框坐标信息绘制矩形检测框。若有标签信息，还会计算合适的字体大小并在检测框上方绘制填充矩形作为背景，再将标签文本绘制在上面。
_make_grid函数：
功能：创建一个二维网格坐标数组，用于后续对模型输出坐标的矫正等操作。
实现：通过np.meshgrid函数分别生成在 x 轴和 y 轴方向的坐标数组，然后将它们堆叠并重塑为形状为(-1, 2)的二维数组，每个元素表示一个网格点的坐标(x, y)。
cal_outputs函数：
功能：对模型输出的坐标等信息进行矫正处理，使其符合实际图像的尺寸和坐标体系。
实现：通过遍历模型的检测层数，根据每层的步长、锚框数量等信息，结合之前创建的网格坐标数组，对模型输出的坐标数据进行一系列的计算和调整，最终返回矫正后的模型输出数据。
post_process_opencv函数：
功能：对模型输出进行后处理，包括计算检测框的实际坐标、进行非极大值抑制（NMS）等操作。
实现：首先从模型输出数据中提取出置信度、中心点坐标、宽度和高度等信息，并根据模型输入和原始输入图像的尺寸关系计算出检测框的实际坐标。然后通过cv2.dnn.NMSBoxes函数进行非极大值抑制操作，根据设定的置信度阈值和非极大值抑制阈值筛选出有效的检测框，最后返回经过处理后的检测框坐标、置信度和类别索引信息。
infer_img函数：
功能：对输入图像进行推理计算，包括图像预处理、模型推理、输出坐标矫正、检测框计算等一系列操作，最终得到检测结果。
实现：
图像预处理：将原始输入图像调整到模型输入的尺寸，转换颜色通道为 RGB，并进行归一化处理，最后将其转换为适合模型输入的张量格式（blob）。
模型推理：使用模型推理会话对象net对预处理后的图像张量进行推理，得到模型的原始输出数据。
输出坐标矫正：调用cal_outputs函数对模型原始输出数据进行坐标矫正。
检测框计算：调用post_process_opencv函数对矫正后的输出数据进行检测框计算，得到最终的检测框坐标、置信度和类别索引信息，并返回这些结果。

def plot_one_box(x, img, color=None, label=None, line_thickness=None) :
    tl = (
            line_thickness or round( 0.002 * (img.shape[0] + img.shape[1]) / 2 ) + 1
    )  # line/font thickness
    color = color or [random.randint(0, 255 ) for _ in range( 3 )]
    c1, c2 = (int( x[0] ), int( x[1] )), (int( x[2] ), int( x[3] ))
    cv2.rectangle( img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA )
    if label :
        tf = max( tl - 1, 1 )  # font thickness
        t_size = cv2.getTextSize( label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle( img, c1, c2, color, -1, cv2.LINE_AA)  # filled
        cv2.putText(
            img,
            label,
            (c1[0], c1[1] - 2),
            0