YoloV8的目标检测推理

YoloV8的目标检测推理

原始的YoloV8封装的层次太高,想要为我们所用可能需要阅读很多API,下面给出比较简单的使用方式

导入所需的库

  • os:用于操作文件系统。
  • cv2 (OpenCV):用于图像处理。
  • numpy:提供数学运算,特别是对数组的操作。
  • ultralytics.YOLO:一个现成的YOLO模型实现,用于对象检测。
  • torch:PyTorch深度学习框架,用于处理模型。
  • serial:用于串口通信。
  • time:用于时间相关的操作。

初始化串口通信的函数:init_serial

  • init_serial 函数用于初始化串口通信。
  • 参数包括端口名称、波特率、字节大小、奇偶校验和停止位。
  • 在尝试打开串口时使用 try-except 结构来捕获异常。
  • 如果串口打开成功,则返回串口对象,否则打印错误信息并退出程序。

加载模型的函数:load_model

  • load_model 函数用于加载YOLO模型。
  • 检查模型权重文件是否存在,不存在则打印错误信息并退出。
  • 使用 YOLO 类从给定路径加载模型,并将其传输到指定的计算设备(CPU或GPU)。
  • 返回加载的模型。

处理图像的函数:process_images

  • process_images 函数负责处理指定路径下的图像。
  • 检查图像路径是否存在,不存在则打印错误信息并退出。
  • 遍历指定路径下的所有 .jpg 图像。
  • 对每张图像应用YOLO模型进行对象检测。
  • 从检测结果中提取类别和边界框信息。
  • 对于检测到的每个对象,绘制边界框并根据类别确定边界框颜色。
  • 根据检测结果通过串口发送"OK"或"not OK"信息。
  • 使用OpenCV显示结果图像,按 q 退出循环。
  • 最后关闭所有OpenCV窗口。

主函数:main

  • 在主函数中设置计算设备。
  • 初始化串口。
  • 加载模型。
  • 调用 process_images 函数处理图像。
  • 使用 if __name__ == "__main__": 确保在直接运行脚本时执行主函数。
python 复制代码
import os
import cv2
import numpy as np
from ultralytics import YOLO
import torch
import serial
import time

def init_serial(port, baudrate, bytesize, parity, stopbits):
    try:
        ser = serial.Serial(port, baudrate, bytesize, parity, stopbits)
        time.sleep(1)  # 等待串口初始化
        return ser
    except serial.SerialException as e:
        print(f"Error opening serial port: {e}")
        exit()

def load_model(weights_path, device):
    if not os.path.exists(weights_path):
        print("Model weights not found!")
        exit()
    model = YOLO(weights_path).to(device)
    model.fuse()
    model.info(verbose=False)
    return model

def process_images(path, model, serial_port):
    if not os.path.exists(path):
        print(f"Path {path} does not exist!")
        exit()

    for img_file in os.listdir(path):
        if not img_file.endswith(".jpg"):
            continue

        img_path = os.path.join(path, img_file)
        img = cv2.imread(img_path)
        if img is None:
            print(f"Failed to load image {img_path}")
            continue

        mask = img.copy()
        result = model(img)
        cls, xywh = result[0].boxes.cls, result[0].boxes.xywh
        cls_, xywh_ = cls.detach().cpu().numpy(), xywh.detach().cpu().numpy()

        for pos, cls_value in zip(xywh_, cls_):
            pt1, pt2 = (np.int_([pos[0] - pos[2] / 2, pos[1] - pos[3] / 2]),
                        np.int_([pos[0] + pos[2] / 2, pos[1] + pos[3] / 2]))
            color = [0, 0, 255] if cls_value == 0 else [0, 255, 0]
            cv2.rectangle(mask, tuple(pt1), tuple(pt2), color, 2)

        res_ = "Yes" if np.any(cls_ == 1) else "No"
        print(res_)
        serial_port.write((res_ + "\r\n").encode())

        cv2.imshow("result", mask)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cv2.destroyAllWindows()

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Using device:", device)

    serial_port = init_serial("/dev/ttyTHS1", 115200, serial.EIGHTBITS, serial.PARITY_NONE, serial.STOPBITS_ONE)
    model = load_model("./weights/best.pt", device)
    process_images("./datasets/pre/", model, serial_port)

if __name__ == "__main__":
    main()
相关推荐
峙峙峙2 分钟前
线性代数--AI数学基础复习
人工智能·线性代数
weiwuxian8 分钟前
揭开智能体的神秘面纱:原来你不是"超级AI"!
人工智能
Codebee8 分钟前
“自举开发“范式:OneCode如何用低代码重构自身工具链
java·人工智能·架构
说私域20 分钟前
基于开源AI智能名片链动2+1模式的S2B2C商城小程序:门店私域流量与视频号直播融合的生态创新研究
人工智能·小程序·开源
Ronin-Lotus23 分钟前
深度学习篇---Yolov系列
人工智能·深度学习
静心问道1 小时前
GoT:超越思维链:语言模型中的有效思维图推理
人工智能·计算机视觉·语言模型
aneasystone本尊1 小时前
学习 Claude Code 的工具使用(三)
人工智能
szxinmai主板定制专家1 小时前
【精密测量】基于ARM+FPGA的多路光栅信号采集方案
服务器·arm开发·人工智能·嵌入式硬件·fpga开发
T__TIII1 小时前
Dify 自定义插件
人工智能·github
快起来别睡了2 小时前
LangChain 介绍及使用指南:从“会聊天”到“能干活”的 AI 应用开发工具
人工智能