基于YOLO实例分割的优质苹果自动分拣装置 简单

pip install ultralytics
pip install onnx==1.17.0
pip install onnxruntime==1.19.2

dataset/                 // 数据集根目录，用于存放训练和验证图像及标签
├── images/              // 图像文件夹，存放原始图片
│   ├── train/           // 训练集图像（用于模型学习）
│   └── val/             // 验证集图像（用于评估模型效果）
│
├── labels/              // 标签文件夹，存放与图像对应的分割掩码标注信息（.txt 或 .json）
│   ├── train/           // 训练集图像的分割标签（polygon 或 YOLO格式）
│   └── val/             // 验证集图像的分割标签
│
└── dataset.yaml         // 数据集配置文件，包含以下信息：
                         // - path：数据集根目录路径（可为绝对或相对路径）
                         // - train：训练图像路径（如 images/train）
                         // - val：验证图像路径（如 images/val）
                         // - nc：分类数量（例如2类苹果）
                         // - names：类别名称列表（如 ['irregular apple', 'pleasing shape apple']）

# train_yolov8_segment.py

from ultralytics import YOLO

# === 第一步：加载模型 ===
# 使用YOLOv8官方提供的预训练模型作为基础，可以选用以下之一：
# yolov8n-seg.pt（Nano，小且快）、yolov8s-seg.pt（Small）、yolov8m-seg.pt（Medium）
model = YOLO("yolov8n-seg.pt")

# === 第二步：开始训练 ===
model.train(
    data="dataset.yaml",      # 指向你的 dataset.yaml 配置文件
    task="segment",           # 明确任务类型为 "segment"
    imgsz=640,                # 图像输入尺寸（推荐为640或更大）
    epochs=20,               # 训练轮次
    batch=16,                 # 每批次处理图像数，受限于显存
    name="apple_seg_model",   # 训练结果输出文件夹名：runs/segment/apple_seg_model
    save=True,                # 是否保存模型权重
    save_period=10,           # 每隔多少epoch保存一次
    val=True,                 # 每个epoch结束后进行验证
    workers=2                 # 数据加载线程数量，适当调小防止卡顿
)

# === 第三步：可选 - 评估最优模型 ===
# 训练完成后自动生成 best.pt，可使用以下方式验证其效果：
# model.val() 会评估 metrics，如 mAP、IoU 等
metrics = model.val()

import torch
from ultralytics import YOLO

# === 第一步：加载训练好的 YOLOv8 模型 ===
# 加载训练好的 best.pt 模型
model = YOLO('runs/segment/apple_seg_model/weights/best.pt')

# === 第二步：设置输入尺寸 ===
input_size = 640  # 设置输入尺寸为 640

# === 第三步：转换为 ONNX 格式 ===
onnx_model_path = 'best.onnx'

# 使用 YOLO 的 export() 方法将模型转换为 ONNX 格式
model.export(format='onnx', imgsz=input_size, dynamic=True)

print(f"✅ 模型已成功转换为 ONNX 格式并保存为 {onnx_model_path}")

import cv2
import numpy as np
from ultralytics import YOLO

# 加载训练好的YOLOv8 Segment模型
model = YOLO("runs/segment/apple_seg_model/weights/best.onnx")

# 读取待推理的图片
image_path = "test1.png"
img = cv2.imread(image_path)

# 推理：使用YOLO模型进行推理
results = model(img)

# 获取第一个推理结果
result = results[0]

# 获取边界框、掩码、类别名映射
boxes = result.boxes
masks = result.masks
names = result.names  # 类别索引映射

# 创建原图副本用于绘图
overlay = img.copy()

# 1. 绘制边界框和标签
for i, box in enumerate(boxes):
    x1, y1, x2, y2 = box.xyxy[0].numpy().astype(int)
    class_id = int(box.cls[0].item())
    class_name = names[class_id]

    # 绘制边框
    cv2.rectangle(overlay, (x1, y1), (x2, y2), (0, 255, 0), 2)

    # 类别标签
    label = f'{class_name}'
    cv2.putText(overlay, label, (x1, y1 - 10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

# 2. 绘制分割掩码
if masks is not None:
    # 获取掩码张量数据，并转为 NumPy
    mask_data = masks.data.cpu().numpy()  # [N, H, W]

    for i in range(mask_data.shape[0]):
        mask_np = mask_data[i]

        # 缩放掩码到原图尺寸
        mask_resized = cv2.resize(mask_np, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_NEAREST)

        # 转换为布尔掩码
        mask_bool = mask_resized > 0.5

        # 设置颜色和透明度
        color = np.array([0, 255, 0], dtype=np.uint8)
        alpha = 0.5

        # 叠加透明掩码颜色
        overlay[mask_bool] = (overlay[mask_bool] * (1 - alpha) + color * alpha).astype(np.uint8)

# 显示结果图像
cv2.imshow('Detected Image with Masks', overlay)
cv2.waitKey(0)
cv2.destroyAllWindows()

# -*- coding: utf-8 -*-
import siot
import cv2
import base64
import numpy as np
import threading
import time
from queue import Queue
from ultralytics import YOLO

# 1. 加载YOLOv8分割模型
model = YOLO("runs/segment/apple_seg_model/weights/best.onnx", task="segment")

# 2. 定义类别颜色映射
class_colors = {
    "irregular apple": (0, 0, 255),        # 红色
    "pleasing shape apple": (0, 255, 0)    # 绿色
}

# 3. MQTT连接配置
client_id = "yolo_pc"
server = "10.1.2.3"
port = 1883
user = "siot"
password = "dfrobot"
topic = "siot/摄像头"

# 4. 创建图像帧队列
frame_queue = Queue(maxsize=5)

# 5. 回调函数：接收图像、推理、入队显示
def on_image_received(client, userdata, msg):
    try:
        data = msg.payload.decode()
        if "base64," not in data:
            return
        b64 = data.split("base64,")[-1]
        img_bytes = base64.b64decode(b64)
        img_array = np.frombuffer(img_bytes, np.uint8)
        frame = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

        if frame is None:
            return

        # ✅ YOLO 推理
        results = model(frame)
        result = results[0]
        boxes = result.boxes
        masks = result.masks
        names = result.names
        overlay = frame.copy()

        # 绘制边界框和标签
        for i, box in enumerate(boxes):
            x1, y1, x2, y2 = box.xyxy[0].numpy().astype(int)
            class_id = int(box.cls[0].item())
            class_name = names[class_id]
            color = class_colors.get(class_name, (255, 255, 255))
            cv2.rectangle(overlay, (x1, y1), (x2, y2), color, 2)
            cv2.putText(overlay, class_name, (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)

        # 绘制分割掩码
        if masks is not None:
            mask_data = masks.data.cpu().numpy()
            for i in range(mask_data.shape[0]):
                mask_np = mask_data[i]
                mask_resized = cv2.resize(mask_np, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST)
                mask_bool = mask_resized > 0.5
                class_id = int(boxes[i].cls[0].item())
                class_name = names[class_id]
                color = np.array(class_colors.get(class_name, (255, 255, 255)), dtype=np.uint8)
                alpha = 0.4
                overlay[mask_bool] = (overlay[mask_bool] * (1 - alpha) + color * alpha).astype(np.uint8)

        # 将推理后的图像放入队列
        if not frame_queue.full():
            frame_queue.put(overlay)
        else:
            frame_queue.get_nowait()
            frame_queue.put(overlay)

        print("[推理] 图像处理完成")

    except Exception as e:
        print(f"[错误] 图像处理失败: {e}")

# 6. 初始化 MQTT
siot.init(client_id=client_id, server=server, port=port, user=user, password=password)
siot.set_callback(on_image_received)
siot.connect()

# 7. 启动 MQTT 后台监听线程
threading.Thread(target=siot.loop, daemon=True).start()
siot.getsubscribe(topic=topic)

# 8. 显示图像主循环
print("[系统] 启动成功，等待图像帧... 按 q 退出")
while True:
    if not frame_queue.empty():
        result_frame = frame_queue.get()
        cv2.imshow("YOLO Segmentation from M10", result_frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 9. 清理资源
cv2.destroyAllWindows()

【输入】：摄像头采集图像（由M10上传）
【推理】：YOLO segment 模型识别每个苹果的类别与轮廓掩码
【可视化】：显示边框、标签和分割结果（带颜色叠加）
【输出】：显示窗口中的图像 + 控制台推理结果提示