首页 > 分享 > 如何使用YOLOv8进行玉米虫害分类的数据集训练玉米虫害分类数据集，包括13个类别共14015张图像

如何使用YOLOv8进行玉米虫害分类的数据集训练玉米虫害分类数据集，包括13个类别共14015张图像

花匠小妙招
2025-04-05 12:52

在这里插入图片描述
玉米虫害分类数据集，包括13个类别共14015张图像：蛴螬（grub）、蝼蛄（mole cricket）、金针虫（wireworm）、白边蛾（white margined moth）、小地老虎（black cutworm）、大地老虎（large cutworm）、黄地老虎（yellow cutworm）、红蜘蛛（red spider）、玉米螟（corn borer）、夜蛾（army worm）、蚜虫（aphids）、白星花金龟（Potosiabre vitarsis）、桃蛀螟（peach borer），训练集、验证集、测试集分别有8404、1399、4212张在这里插入图片描述

如何使用YOLOv8进行玉米虫害分类的数据集训练。由于这是一个分类任务而不是目标检测任务，我们将使用YOLOv8的分类功能来进行训练。

使用 YOLOv8 训练玉米虫害分类

在这里插入图片描述

数据集信息类别: 13类 (蛴螬(grub), 蝼蛄(mole cricket), 金针虫(wireworm), 白边蛾(white margined moth), 小地老虎(black cutworm), 大地老虎(large cutworm), 黄地老虎(yellow cutworm), 红蜘蛛(red spider), 玉米螟(corn borer), 夜蛾(army worm), 蚜虫(aphids), 白星花金龟(Potosiabre vitarsis), 桃蛀螟(peach borer))图片数量: 共14015张分辨率: 假设为统一的分辨率（例如 224x224）标签格式: 分类格式 (txt 或 CSV) 数据集划分: 训练集: 8404 张验证集: 1399 张测试集: 4212 张步骤概述 数据集准备创建数据集配置文件 (data.yaml)训练模型评估模型可视化训练结果清理临时文件推理和显示结果 详细步骤 1. 数据集准备

确保你的数据集已经按照上述格式准备好，并且包含 train, val, 和 test 目录。

corn_pest_classification/ ├── datasets/ │ └── corn_pest_dataset/ │ ├── train/ │ │ ├── images/ │ │ │ ├── image1.jpg │ │ │ ├── image2.jpg │ │ │ └── ... │ │ └── labels/ │ │ ├── image1.txt │ │ ├── image2.txt │ │ └── ... │ ├── val/ │ │ ├── images/ │ │ │ ├── image1.jpg │ │ │ ├── image2.jpg │ │ │ └── ... │ │ └── labels/ │ │ ├── image1.txt │ │ ├── image2.txt │ │ └── ... │ └── test/ │ ├── images/ │ │ ├── image1.jpg │ │ ├── image2.jpg │ │ └── ... │ └── labels/ │ ├── image1.txt │ ├── image2.txt │ └── ... └── main.py 12345678910111213141516171819202122232425262728293031 2. 创建数据集配置文件 (data.yaml)

创建一个 data.yaml 文件来配置数据集路径和类别信息。

train: ./datasets/corn_pest_dataset/train/images val: ./datasets/corn_pest_dataset/val/images nc: 13 # 类别数量 names: ['grub', 'mole_cricket', 'wireworm', 'white_margined_moth', 'black_cutworm', 'large_cutworm', 'yellow_cutworm', 'red_spider', 'corn_borer', 'army_worm', 'aphids', 'Potosiabre_vitarsis', 'peach_borer'] 12345 3. 训练模型

使用YOLOv8进行分类训练。

import torch from ultralytics import YOLO # 设置随机种子以保证可重复性 torch.manual_seed(42) # 定义数据集路径 dataset_config = 'data.yaml' # 加载预训练的YOLOv8n-c模型（用于分类） model = YOLO('yolov8n-cls.pt') # 训练模型 results = model.train( data=dataset_config, epochs=100, imgsz=224, batch=16, name='corn_pest_classification', project='runs/classify' ) # 评估模型 metrics = model.val() # 保存最佳模型权重 best_model_weights = 'runs/classify/corn_pest_classification/weights/best.pt' print(f"Best model weights saved to {best_model_weights}") 12345678910111213141516171819202122232425262728 4. 可视化训练结果

可视化训练结果。

from ultralytics import YOLO # 加载训练好的模型 model = YOLO('runs/classify/corn_pest_classification/weights/best.pt') # 可视化训练结果 model.plot_results(save=True, save_dir='runs/classify/corn_pest_classification') 1234567 5. 清理临时文件

清理不必要的临时文件。

import shutil def clean_temp_files(project_dir): temp_dirs = [ f'{project_dir}/wandb', f'{project_dir}/cache' ] for dir_path in temp_dirs: if os.path.exists(dir_path): shutil.rmtree(dir_path) print(f"Removed directory: {dir_path}") # 使用示例 clean_temp_files('runs/classify/corn_pest_classification') 123456789101112131415 6. 推理和显示结果

推理和显示结果。

from ultralytics import YOLO import cv2 import numpy as np from PIL import Image import matplotlib.pyplot as plt # 主函数 def main(): # 加载模型 model = YOLO('runs/classify/corn_pest_classification/weights/best.pt') # 测试图片路径 test_image_path = './datasets/corn_pest_dataset/test/images/test_image.jpg' # 进行预测 results = model.predict(test_image_path, conf=0.5)[0] # 获取预测结果 predicted_class_id = int(results.probs.argmax()) predicted_class_name = model.names[predicted_class_id] confidence = float(results.probs.max()) # 打印预测结果 print(f"Predicted Class: {predicted_class_name}, Confidence: {confidence:.2f}") # 显示图片 image = cv2.imread(test_image_path) image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 在图片上添加预测标签 font = cv2.FONT_HERSHEY_SIMPLEX text = f"{predicted_class_name} ({confidence:.2f})" position = (10, 30) font_scale = 1 color = (0, 255, 0) thickness = 2 cv2.putText(image_rgb, text, position, font, font_scale, color, thickness) plt.imshow(image_rgb) plt.axis('off') plt.show() if __name__ == "__main__": main() 123456789101112131415161718192021222324252627282930313233343536373839404142434445 运行脚本

在终端中运行以下命令来执行整个流程：

python main.py 1 总结

以上文档包含了从数据集准备、模型训练、评估、可视化训练结果、清理临时文件到推理和显示结果的所有步骤。希望这些详细的信息和代码能够帮助你顺利实施和优化你的玉米虫害分类系统。如果有任何进一步的问题或需要更多帮助，请随时提问！

这段代码涵盖了从数据集准备、模型训练、评估、可视化训练结果、清理临时文件到推理和显示结果的所有步骤。你可以根据需要调整参数和路径。