如何使用YOLOv8来训练——9200张图像毒蘑菇检测数据集-CSDN博客

本文链接：https://blog.csdn.net/2401_88440984/article/details/143355442

如何使用YOLOv8来训练一个包含9200张图像的毒蘑菇检测数据集，并附上详细的训练代码和步骤。

数据集描述

该数据集包含以下信息：

数据量：9200张图像
类别：14类
- Amanita citrina
- Gyromitra infula
- Hygrophoropsis aurantiaca
- Imleria badia
- Lactarius turpis
- Boletus reticulatus
- Amanita pantherina
- Coltricia perennis
- Amanita hemibapha
- Amanita javanica
- Amanita princeps
- Russula emetica
- （待补充两个类别名称）

数据集组织

假设你的数据集目录结构如下：

toxic_mushroom_dataset/
├── train/
│   ├── images/
│   └── labels/
├── valid/
│   ├── images/
│   └── labels/
└── data.yaml  # 数据配置文件

其中：

train/ 包含训练集的图像和TXT标签文件。
valid/ 包含验证集的图像和TXT标签文件。
data.yaml 包含数据配置信息。

数据配置文件

创建或确认data.yaml文件是否正确配置了数据集路径和类别信息：

train: ./train/images/  # 训练集图像路径
val: ./valid/images/  # 验证集图像路径
test: ./test/images/  # 测试集图像路径（如果有）

# Classes
nc: 14  # 类别数量
names: ['amanita_citrina', 'gyromitra_infula', 'hygrophoropsis_aurantiaca', 
        'imleria_badia', 'lactarius_turpis', 'boletus_reticulatus',
        'amanita_pantherina', 'coltricia_perennis', 'amanita_hemibapha',
        'amanita_javanica', 'amanita_princeps', 'russula_emetica', 
        'missing_category_1', 'missing_category_2']  # 类别名称列表

VOC格式转YOLO格式

如果你的数据集中有VOC格式的标签文件，需要将其转换为YOLO格式的TXT文件。这里提供一个Python脚本示例

import os
import xml.etree.ElementTree as ET
from PIL import Image

def convert_voc_to_yolo(xml_file, img_size, class_names, output_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    
    width = int(root.find('size/width').text)
    height = int(root.find('size/height').text)
    
    labels = []
    for obj in root.iter('object'):
        cls = obj.find('name').text.lower()  # 将类别名转换为小写
        if cls not in class_names:
            continue
        cls_id = class_names.index(cls)
        
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),
             float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        
        bb = convert_coordinates(b, width, height, img_size)
        labels.append([cls_id, *bb])
    
    with open(output_file, 'w') as f:
        for label in labels:
            f.write('%d %.6f %.6f %.6f %.6f\n' % tuple(label))

def convert_coordinates(box, orig_w, orig_h, target_size):
    dw, dh = target_size
    x_center = ((box[1] + box[0]) / 2) / orig_w
    y_center = ((box[3] + box[2]) / 2) / orig_h
    w = (box[1] - box[0]) / orig_w
    h = (box[3] - box[2]) / orig_h
    
    x_center *= dw
    y_center *= dh
    w *= dw
    h *= dh
    
    return x_center, y_center, w, h

def main():
    class_names = ['amanita_citrina', 'gyromitra_infula', 'hygrophoropsis_aurantiaca', 
                   'imleria_badia', 'lactarius_turpis', 'boletus_reticulatus',
                   'amanita_pantherina', 'coltricia_perennis', 'amanita_hemibapha',
                   'amanita_javanica', 'amanita_princeps', 'russula_emetica', 
                   'missing_category_1', 'missing_category_2']  # 替换为实际类别名称
    img_size = (640, 640)  # 输入图像大小
    
    for phase in ['train', 'valid']:
        xml_dir = f'./toxic_mushroom_dataset/{phase}/annotations/'
        txt_dir = f'./toxic_mushroom_dataset/{phase}/labels_txt/'
        os.makedirs(txt_dir, exist_ok=True)
        
        for xml_file in os.listdir(xml_dir):
            if xml_file.endswith('.xml'):
                img_file = xml_file.replace('.xml', '.jpg')
                img_path = os.path.join(f'./toxic_mushroom_dataset/{phase}/images/', img_file)
                img = Image.open(img_path)
                img_width, img_height = img.size
                
                output_file = os.path.join(txt_dir, img_file.replace('.jpg', '.txt'))
                convert_voc_to_yolo(os.path.join(xml_dir, xml_file), img_size, class_names, output_file)

if __name__ == '__main__':
    main()

运行上述脚本后，会在toxic_mushroom_dataset/train/labels_txt/和toxic_mushroom_dataset/valid/labels_txt/目录下生成YOLO格式的.txt标签文件。

安装YOLOv8

如果你还没有安装YOLOv8，可以使用以下命令安装

pip install ultralytics

训练模型

使用YOLOv8训练模型的命令非常简单，你可以直接使用以下命令开始训练：

cd path/to/toxic_mushroom_dataset/

# 下载预训练权重
wget https://github.com/ultralytics/ultralytics/releases/download/v8.0.19/yolov8n.pt

# 开始训练
ultralytics train model=yolov8n.yaml data=./data.yaml epochs=100 imgsz=640

在这个命令中：

model=yolov8n.yaml：指定使用的YOLOv8模型配置文件。
data=./data.yaml：指定数据配置文件。
epochs=100：训练轮数。
imgsz=640：输入图像的大小。

使用预定义配置

YOLOv8提供了多种预定义模型配置，例如yolov8n, yolov8s, yolov8m, yolov8l, yolov8x。你可以根据需求选择不同的模型配置进行训练。例如：

ultralytics train model=yolov8s.yaml data=./data.yaml epochs=100 imgsz=640

模型评估

训练完成后，可以使用以下命令评估模型在验证集上的表现：

ultralytics val model=best.pt data=./data.yaml imgsz=640

这里的best.pt是训练过程中产生的最佳模型权重文件。

模型预测

你可以使用训练好的模型对新图像进行预测：

ultralytics predict model=best.pt source=path/to/your/image.jpg imgsz=640

查看训练结果

训练过程中的日志和结果会保存在runs/detect/目录下，你可以查看训练过程中的损失、精度等信息。

注意事项

数据集质量：确保数据集的质量，包括清晰度、标注准确性等。
模型选择：可以选择更强大的模型版本（如YOLOv8m、YOLOv8l等）以提高性能。
超参数调整：根据实际情况调整超参数，如批量大小（batch-size）、图像大小（imgsz）等。
监控性能：训练过程中监控损失函数和mAP指标，确保模型收敛。

通过上述步骤，你可以使用YOLOv8来训练一个毒蘑菇检测数据集，并使用训练好的模型进行预测