基于YOLO的水果目标检测装置 简单

pip install ultralytics
pip install onnx==1.16.1
pip install onnxruntime==1.17.1

pip install onnx==1.16.1
pip install onnxruntime==1.17.1
pip install opencv-python
pip install pinpong

pip list

fruits_dataset/             //水果数据集
│  
├── train/                  // 训练集，使用训练集的数据进行模型训练
│   ├── images/             // 训练集的图片文件夹
│   └── labels/             // 训练集中图片的标注信息文件夹，作用是提供图像中目标的位置和类别信息。
│  
├── val/                    // 验证集，验证集用于评估模型在未见过的数据上的表现
│   ├── images/             // 验证集的图片文件夹
│   └── labels/             // 验证集图片的标注信息文件夹
│  
└── data.yaml             // 数据集配置文件，用于定义数据集的参数，通常是一个 YAML 文件。

from ultralytics import YOLO  # 导入YOLO类，用于加载和训练模型
import time  # 导入time模块，用于延迟操作
import os  # 导入os模块，用于操作文件路径

# 打印当前工作目录，帮助确认代码运行的路径是否正确
#print("Current working directory:", os.getcwd())

# 确保路径正确
model = YOLO('yolov8s.pt')  # 加载YOLOv8s预训练模型,yolov8s.pt是一个轻量级模型，适合资源有限的环境

# 开始训练模型
results = model.train(
    data=os.path.join(os.getcwd(), '..', 'yamls', 'data.yaml'),  # 数据配置文件路径
    # 使用相对路径，确保指向正确的yamls文件夹,data.yaml 文件中定义了数据集的路径、类别数等信息
    epochs=2,  # 设置训练轮数为2（通常需要更多轮数，这里可能是测试）
    imgsz=320,  # 设置输入图像大小为320x320，降低图像大小可以减少显存占用，适合资源有限的环境
    device='cpu',  # 使用CPU进行训练（如果无GPU或CUDA环境）
    # 如果有GPU，可以设置为 '0' 或 'cuda'
    workers=0,  # 数据加载线程数（Windows系统建议设置为0）
    # workers=0 表示主进程加载数据，避免多线程问题
    batch=2,  # 设置批量大小为2，降低批量大小可以减少显存占用，但可能影响训练效率
    cache=False  # 不使用缓存（避免内存占用过多）
)

from ultralytics import YOLO

# 模型路径，这里填你的pt文件的实际路径
model_path = r'C:\Users\li720\Downloads\ultralytics-main\ultralytics-main\runs\runs\detect\train5\weights\best.pt'  # 训练完成后保存的模型路径

# 加载模型
model = YOLO(model_path)

# 导出为onnx格式
model.export(format='onnx')

print("模型已成功导出为onnx格式")

from ultralytics import YOLO
import cv2

# 1. 加载预训练模型
model = YOLO('best.onnx', task='detect')  # 显式指定任务为检测

# 2. 进行图像推理（带类别统计）
def predict_image(image_path, conf=0.5):
    # 读取图像并调整尺寸
    img = cv2.imread(image_path)
    img = cv2.resize(img, (320, 320))
    
    # 执行推理
    results = model.predict(
        source=img,
        conf=conf,
        save=False,
        imgsz=320
    )
    
    # 初始化类别计数器
    class_counts = {}
    
    # 遍历检测结果
    for result in results:
        # 获取检测到的所有类别
        detected_classes = result.boxes.cls.tolist()
        
        # 统计每个类别的出现次数
        for class_id in detected_classes:
            class_name = result.names[int(class_id)]
            class_counts[class_name] = class_counts.get(class_name, 0) + 1
    
    # 可视化结果
    annotated_img = results[0].plot()
    
    # 显示结果
    cv2.imshow('Detection', annotated_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    
    return results, class_counts  # 返回结果和统计信息

if __name__ == '__main__':
    try:
        # 执行检测并获取统计结果
        image_results, counts = predict_image('test.png', conf=0.6)
        
        # 打印检测统计
        print("\n检测结果统计:")
        for class_name, count in counts.items():
            print(f"{class_name}: {count}个")
            
        # 原始检测信息（可选）
        print("\n详细检测信息:")
        for result in image_results:
            for box in result.boxes:
                print(f"类别: {result.names[int(box.cls)]}, 置信度: {box.conf.item():.2f}")
                
    except Exception as e:
        print(f"发生错误: {str(e)}")

import cv2  # 导入OpenCV库，用于图像处理和视频捕获
from ultralytics import YOLO  # 导入YOLO模型库，用于目标检测

# 加载YOLO模型
model = YOLO('best.onnx', task='detect')  # 加载模型文件 'best.onnx'，并明确指定任务类型为检测任务

# 初始化摄像头
cap = cv2.VideoCapture(0)  # 打开默认摄像头（设备编号为0）
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320)  # 设置摄像头捕获的帧宽度为320像素
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 320)  # 设置摄像头捕获的帧高度为320像素

# 检查摄像头是否成功打开
if not cap.isOpened():
    print("无法打开摄像头，请检查设备编号或摄像头连接。")
    exit()  # 如果摄像头无法打开，程序退出

try:
    # 进入主循环，实时处理视频流
    while True:
        ret, frame = cap.read()  # 从摄像头读取一帧图像
        if not ret:  # 如果读取失败（例如摄像头关闭），退出循环
            break
        
        # 转换颜色空间，将BGR格式的图像转换为RGB格式
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # YOLO模型需要RGB格式的输入
        
        # 使用YOLO模型进行目标检测
        results = model(rgb_frame, imgsz=320, verbose=False, conf=0.2)  # 对RGB图像进行推理，指定输入图像尺寸为320x320，并设置置信度阈值0.2才输出
        
        # 获取检测结果
        detections = results[0].boxes.data.tolist()  # 提取检测框信息，将其转换为列表格式
        
        # 统计水果数量
        fruit_counts = {}  # 创建一个字典用于存储每种水果的数量
        for detection in detections:  # 遍历每个检测结果
            class_id = int(detection[5])  # 获取类别ID（位于检测框数据的第6个位置）
            class_name = model.names[class_id]  # 根据类别ID获取对应的类别名称（例如水果名称）
            
            # 更新水果计数
            fruit_counts[class_name] = fruit_counts.get(class_name, 0) + 1  # 如果类别已存在，计数加1；否则初始化为1
        
        # 绘制检测结果
        annotated_frame = results[0].plot()  # 使用YOLO库自带的功能绘制检测框和类别标签
        
        # 在图像上添加水果数量统计信息
        y_offset = 30  # 初始化文本的垂直偏移量
        for fruit, count in fruit_counts.items():  # 遍历统计结果
            cv2.putText(  # 在图像上绘制文本
                annotated_frame, 
                f"{fruit}: {count}",  # 文本内容，显示水果名称和数量
                (10, y_offset),  # 文本起始位置
                cv2.FONT_HERSHEY_SIMPLEX,  # 字体类型
                0.7,  # 字体大小
                (0, 255, 0),  # 文本颜色（绿色）
                2  # 文本线条粗细
            )
            y_offset += 30  # 更新文本的垂直偏移量，避免文本重叠
        
        # 显示结果图像
        cv2.imshow(  # 显示处理后的图像
            'Detection',  # 窗口标题
            cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)  # 将RGB格式的图像转换为BGR格式
        )
        
        # 按下 'a' 键退出程序
        if cv2.waitKey(1) == ord('a'):  # 检测按键事件，如果按下 'a' 键
            break

except KeyboardInterrupt:  # 捕获键盘中断（例如用户按下Ctrl+C）
    print("程序被用户中断。")

finally:  # 无论程序是否正常退出，都会执行以下代码
    cap.release()  # 释放摄像头资源
    cv2.destroyAllWindows()  # 关闭所有OpenCV创建的窗口

conda activate yolo

python cv_inference.py

    # 使用YOLO模型进行目标检测
   results = model(rgb_frame, imgsz=320, verbose=False, conf=0.5)  # 对RGB图像进行推理，指定输入图像尺寸为320x320，并设置置信度阈值0.2才输出
    print('模型的原始输出:',results)

        # 获取检测结果
        detections = results[0].boxes.data.tolist()  # 提取检测框信息，将其转换为列表格式
        
        # 统计水果数量
        fruit_counts = {}  # 创建一个字典用于存储每种水果的数量
        for detection in detections:  # 遍历每个检测结果
            class_id = int(detection[5])  # 获取类别ID（位于检测框数据的第6个位置）
            class_name = model.names[class_id]  # 根据类别ID获取对应的类别名称（例如水果名称）
            
            # 更新水果计数
            fruit_counts[class_name] = fruit_counts.get(class_name, 0) + 1  # 如果类别已存在，计数加1；否则初始化为1
        
        # 绘制检测结果
        annotated_frame = results[0].plot()  # 使用YOLO库自带的功能绘制检测框和类别标签
        
        # 在图像上添加水果数量统计信息
        y_offset = 30  # 初始化文本的垂直偏移量
        for fruit, count in fruit_counts.items():  # 遍历统计结果
            cv2.putText(  # 在图像上绘制文本
                annotated_frame, 
                f"{fruit}: {count}",  # 文本内容，显示水果名称和数量
                (10, y_offset),  # 文本起始位置
                cv2.FONT_HERSHEY_SIMPLEX,  # 字体类型
                0.7,  # 字体大小
                (0, 255, 0),  # 文本颜色（绿色）
                2  # 文本线条粗细
            )
            y_offset += 30  # 更新文本的垂直偏移量，避免文本重叠