使用YOLOv11实现水果类别检测:从数据到模型训练的全过程
项目背景
在日常生活中,水果的类别识别直接影响库存管理和质量控制。传统检测依赖人工,效率低下且主观。通过计算机视觉,我们可以自动化这个过程。YOLO系列模型以其速度和准确性闻名,这次我使用YOLOv11n进行训练,目标是检测31种水果类别,包括苹果、香蕉、橙子等。
数据集来源于公开来源,图像标注为YOLO格式,包含训练、验证和测试集。nc=31,names包括’fresh_apple’到’rotten_tomato’等。
方法步骤
- 数据准备:数据集已准备好,包含标注图像。
-
模型训练:使用训练脚本,配置YOLO数据。训练参数:epochs=50, batch=16, imgsz=640。
以下是训练脚本的主要代码:
from ultralytics import YOLO import os import glob import torchdef main():model = YOLO('models/yolo11n.pt')yaml_path = 'data/data.yaml'results = model.train(data=yaml_path, epochs=50, batch=16, imgsz=640, name='yolo11n_fruit', project='fruit_results')print("Training completed. Results saved in:", results.save_dir)if __name__ == '__main__':main()
训练命令:
python train_fire_smoke_tassel.py
结果保存在
fruit_results/yolo11n_fruit_20250716_210401
。
-
模型验证:使用
test_saved_model.py
验证,加载best.pt模型。以下是验证脚本的主要代码:
from ultralytics import YOLO import os import globdef main():yaml_path = 'data/data.yaml'results_dir = 'fruit_results'latest_run = max(glob.glob(os.path.join(results_dir, 'yolo11n_fruit_*')), key=os.path.getmtime)model_path = os.path.join(latest_run, 'weights', 'best.pt')model = YOLO(model_path)metrics = model.val(data=yaml_path)print(f"Validation metrics: {metrics}")if __name__ == '__main__':main()
关键指标:
- mAP50: 0.709
- mAP50-95: 0.421
- Precision: 0.675
- Recall: 0.715
结果分析
模型在验证集上表现不错,mAP50达0.709,表示中等难度下的检测准确。Recall较高,意味着漏检少,适合实际应用。但mAP50-95较低,说明小目标或复杂背景需优化。
预测示例:
未来改进
- 增加数据增强,提升鲁棒性。
- 尝试更大模型如YOLOv11m。
- 部署到移动端,实现实时检测。