土豆瑕疵检测项目实施方案 (基于 YOLOv8)

土豆瑕疵检测项目实施方案 (基于 YOLOv8)

1. 环境准备与工具链选择

环节	工具/语言	作用	备注
编程语言	Python 3.8+	核心开发语言	必须。
深度学习框架	PyTorch 2.0+	YOLOv8 的底层框架	必须。
模型框架	Ultralytics YOLOv8	快速训练和部署对象检测模型	必须。
标注工具	LabelImg / Roboflow / LabelMe	制作 YOLO 格式的标签文件	推荐使用 LabelImg 或 Roboflow。
硬件要求	具备 NVIDIA GPU (推荐 8GB 显存以上)	确保训练效率和速度	必须。

安装核心库:

# 安装 PyTorch (根据你的CUDA版本选择)
# pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 

# 安装 Ultralytics (包含 YOLOv8)
pip install ultralytics

# 其他常用工具
pip install opencv-python Pillow pyyaml

---

2. 数据集采集与准备 (Data Collection)

高质量的数据集是模型成功的基石。

2.1 样本采集策略

类别	描述	目标数量
正样本	包含各种类型瑕疵的土豆图像。	3000 - 5000 张 (初期目标)
负样本	健康、无瑕疵的土豆图像。	约占正样本的 10-20%
背景样本	仅包含传送带、背景或环境光干扰的图像。	约占正样本的 5%

采集要求 (力求达到使用标准):

1. 光照统一性： 尽量在模拟实际流水线的光照条件下采集，避免强烈的反光或阴影。
2. 角度多样性： 采集土豆各个侧面的图像，确保模型能检测到隐藏在侧面的瑕疵。
3. 瑕疵多样性： 确保样本覆盖所有需要检测的瑕疵类型，例如：绿皮 (greening)、疮痂 (scab)、机械损伤 (damage)、虫洞 (wormhole) 等。
4. 图像尺寸： 统一采集高分辨率图像（如 1080P 或更高），但训练前需缩放至 YOLOv8 要求的尺寸（如 640x640）。

2.2 定义瑕疵类别 (Classes Definition)

确定模型需要识别的瑕疵种类。

序号	类别名称 (英文)	类别名称 (中文)
0	damage	机械损伤 (切割、裂纹等)
1	greening	绿皮 (土豆变绿区域)
2	scab	疮痂 (土豆表面粗糙病斑)
3	wormhole	虫洞 (小且深的孔洞)

---

3. 数据标注 (Data Annotation)

使用标注工具对采集的图像进行标注。

1. 标注工具： 使用 LabelImg (推荐)，它能导出标准的 YOLO 格式 的标签文件。
2. 标注标准：
   • 使用矩形框（Bounding Box）紧密地框选每一种瑕疵区域。
   • 对于同一个土豆上的多个不同瑕疵，需要标注多个不同类别的框。
   • 健康土豆的图片作为负样本，无需标注任何框，但其 .txt 标签文件也需保留（内容为空）。
3. 输出格式： 标注完成后，每张图像 image.jpg 对应一个 image.txt 文件，内容格式为：<class_id> <center_x> <center_y> <width> <height> (坐标均为归一化到 0-1 的相对值)。

---

4. 数据集划分与目录结构

将所有图像和标签文件组织成标准的 YOLO 格式。

4.1 目录结构

potato_detection/
├── images/
│   ├── train/  (训练集图像)
│   ├── val/    (验证集图像)
│   └── test/   (测试集图像)
└── labels/
    ├── train/  (训练集标签 *.txt)
    ├── val/    (验证集标签 *.txt)
    └── test/   (测试集标签 *.txt)

4.2 划分比例

推荐使用以下比例进行数据集划分：

集合	作用	推荐比例
Train (训练集)	用于模型学习特征和权重调整。	70% - 80%
Val (验证集)	用于在训练过程中评估模型性能，辅助调整超参数，但不参与权重更新。	10% - 15%
Test (测试集)	用于最终评估模型在未知数据上的泛化能力，训练过程中绝对不能使用。	10% - 15%

4.3 配置文件 (potato_data.yaml)

创建一个 YAML 文件来告诉 YOLOv8 你的数据集位置和类别信息。

# potato_data.yaml
path: /path/to/potato_detection  # 替换为你的项目根目录
train: images/train  # 训练图像路径
val: images/val      # 验证图像路径
test: images/test    # 测试图像路径 (可选，用于最终测试)

# 类别信息
names:
  0: damage
  1: greening
  2: scab
  3: wormhole

# 类别数量
nc: 4 

---

5. 模型训练 (Model Training)

5.1 模型选择

选用 YOLOv8n (Nano) 或 YOLOv8s (Small)。

• YOLOv8n： 速度最快，体积最小，适合边缘设备或对速度要求极高的场景（推荐作为起点）。
• YOLOv8s： 性能稍好，但速度稍慢，如果 n 模型精度不够，可选用 s 模型。

5.2 训练参数配置

参数	推荐值	描述
model	yolov8n.pt	使用预训练的 Nano 模型权重。
data	potato_data.yaml	上一步创建的数据集配置文件。
epochs	100 - 300	训练轮次，根据收敛情况调整。
imgsz	640	输入图像尺寸。
batch	16 - 64	批处理大小，根据 GPU 显存调整。
patience	50	早停机制，如果在 50 个 epoch 内 mAP 没有提升则停止训练。
device	0 或 0,1	指定使用的 GPU 编号。

5.3 训练代码 (Python Script)

from ultralytics import YOLO

# 1. 加载模型
# 建议从预训练的yolov8n模型开始
model = YOLO('yolov8n.pt')  

# 2. 模型训练
# 命令行等价于: yolo train model=yolov8n.pt data=potato_data.yaml epochs=200 imgsz=640 batch=16
results = model.train(
    data='potato_data.yaml',   # 数据集配置文件
    epochs=200,                # 训练轮次
    imgsz=640,                 # 图像尺寸
    batch=16,                  # 批处理大小
    name='potato_detection_v1',# 训练结果保存目录名
    device=0,                  # 使用第一块GPU
    patience=50,               # 50轮无提升则早停
    optimizer='AdamW'          # 推荐使用AdamW优化器
)

print("模型训练完成，结果保存在 runs/detect/potato_detection_v1 目录下。")

---

6. 模型验证与测试

6.1 模型验证 (Validation)

训练过程中，模型会自动在验证集 (val) 上进行验证，输出关键指标。

关键指标：

指标	描述	使用标准要求
Precision §	精度，模型预测正确的比例（查准率）。	≥ 90%
Recall ®	召回率，模型找出的真实目标的比例（查全率）。	≥ 90%
mAP50	mean Average Precision (IoU=0.5)，衡量模型的综合性能。	≥ 90%
mAP50-95	平均 mAP (从 IoU=0.5 到 0.95)，更严格的综合性能衡量。	≥ 85%

6.2 模型测试 (Testing)

训练结束后，使用性能最佳的权重文件（通常是 runs/detect/potato_detection_v1/weights/best.pt）在测试集 (test) 上进行最终评估。

# 3. 模型测试
# 加载最佳权重文件
best_model = YOLO('runs/detect/potato_detection_v1/weights/best.pt')

# 在测试集上运行评估
metrics = best_model.val(
    data='potato_data.yaml', 
    split='test', # 明确指定在测试集上运行
    imgsz=640
)

# 打印最终测试结果
print(f"最终测试 mAP50: {metrics.box.map50}")
print(f"最终测试 mAP50-95: {metrics.box.map}")

---

7. 部署与使用标准

为了使模型达到工业使用标准，需要进行以下优化和部署准备：

7.1 模型优化与导出

在达到满意的精度后，将模型导出为更高效的格式以提高推理速度。

# 4. 模型导出 (转换为ONNX或TensorRT以提高部署速度)
best_model = YOLO('runs/detect/potato_detection_v1/weights/best.pt')

# 导出为 ONNX 格式（兼容性高）
best_model.export(format='onnx', opset=12) 

# 导出为 TensorRT 格式（NVIDIA GPU 上的最高性能）
# best_model.export(format='engine') 

7.2 实时推理 (Inference)

部署阶段的代码将专注于使用导出的模型对实时视频流或图像进行快速检测。

import cv2

# 加载 ONNX 格式的模型（假设已导出）
model_onnx = YOLO('runs/detect/potato_detection_v1/weights/best.onnx')

# 实时视频流检测（示例）
cap = cv2.VideoCapture(0) # 或视频文件路径

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # 运行检测 (Conf: 置信度阈值, IOU: NMS阈值)
    results = model_onnx(frame, conf=0.7, iou=0.45) 

    # 绘制结果
    annotated_frame = results[0].plot()

    cv2.imshow("Potato Defect Detection", annotated_frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

通过遵循上述详细的实施方案，特别是严格控制数据采集质量和精确的标注，您可以高效地训练出一个满足土豆瑕疵检测使用标准的 YOLOv8 模型。