如何使用YOLOv8来训练一个包含926张图像的红外弱小飞机目标检测数据集,并附上详细的训练代码和步骤。

数据集描述

该数据集包含以下信息:

  • 数据量:926张图像
  • 类别:1类(bird)
  • 标注格式:BMP图像和XML标签文件一一对应

数据集组织

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

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

其中:

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

XML转YOLO格式

由于YOLOv8要求标签文件为.txt格式,并且标签内容为归一化后的边界框坐标,因此你需要先将XML格式的标签文件转换为YOLO格式的.txt文件。你可以使用Python脚本来实现这个转换:

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

def convert_xml_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
        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 = ['bird']
    img_size = (640, 640)  # 输入图像大小
    
    for phase in ['train', 'valid']:
        xml_dir = f'./infrared_aircraft_dataset/{phase}/labels_xml/'
        txt_dir = f'./infrared_aircraft_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', '.bmp')
                img_path = os.path.join(f'./infrared_aircraft_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('.bmp', '.txt'))
                convert_xml_to_yolo(os.path.join(xml_dir, xml_file), img_size, class_names, output_file)

if __name__ == '__main__':
    main()

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

数据配置文件

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

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

# Classes
nc: 1  # 类别数量
names: ['bird']  # 类别名称列表

安装YOLOv8

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

pip install ultralytics

训练模型

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

cd path/to/infrared_aircraft_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.bmp imgsz=640

查看训练结果

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

注意事项

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

通过上述步骤,你可以使用YOLOv8来训练一个红外弱小飞机目标检测数据集,并使用训练好的模型进行预测。

更多推荐