如何训练自己的数据集之——DIOR遥感小目标目标检测数据集

 DIOR_dataset/

    
 ├── train/
    
 │   ├── images/
    
 │   │   ├── image_0001.jpg
    
 │   │   ├── image_0002.jpg
    
 │   │   └── ...
    
 │   ├── labels_yolo/
    
 │   │   ├── image_0001.txt
    
 │   │   ├── image_0002.txt
    
 │   │   └── ...
    
 │   ├── labels_voc/
    
 │   │   ├── image_0001.xml
    
 │   │   ├── image_0002.xml
    
 │   │   └── ...
    
 ├── valid/
    
 │   ├── images/
    
 │   │   ├── image_0001.jpg
    
 │   │   ├── image_0002.jpg
    
 │   │   └── ...
    
 │   ├── labels_yolo/
    
 │   │   ├── image_0001.txt
    
 │   │   ├── image_0002.txt
    
 │   │   └── ...
    
 │   ├── labels_voc/
    
 │   │   ├── image_0001.xml
    
 │   │   ├── image_0002.xml
    
 │   │   └── ...
    
 └── data.yaml  # 数据集配置文件
    
    
    
    
    AI写代码

 # data.yaml 文件

    
 train: ../DIOR_dataset/train/images/
    
 val: ../DIOR_dataset/valid/images/
    
 test: null  # 如果有单独的测试集路径，则填写
    
  
    
 nc: 20  # number of classes
    
 names: ['harbor', 'ship', 'storagetank', 'chimney', 'dam', 'trainstation', 'basketballcourt', 'airport', 'expressway-service-area', 'airplane', 'baseballfield', 'expressway-toll-station', 'vehicle', 'golffield', 'bridge', 'groundtrackfield', 'overpass', 'windmill', 'tenniscourt', 'stadium']
    
    
    
    
    AI写代码

 git clone https://github.com/ultralytics/ultralytics.git  # 克隆YOLOv8仓库

    
 cd ultralytics
    
 pip install -r requirements.txt  # 安装依赖项
    
    
    
    
    AI写代码

 import xml.etree.ElementTree as ET

    
 import os
    
  
    
 def convert_voc_to_yolo(xml_file_path, txt_file_path, img_width, img_height):
    
     tree = ET.parse(xml_file_path)
    
     root = tree.getroot()
    
     
    
     with open(txt_file_path, 'w') as txt_file:
    
     for obj in root.iter('object'):
    
         difficult = obj.find('difficult').text
    
         cls = obj.find('name').text
    
         if int(difficult) == 1 or int(difficult) == -1:
    
             continue
    
         cls_id = names.index(cls)
    
         xmlbox = obj.find('bndbox')
    
         b = (int(float(xmlbox.find('xmin').text)),
    
              int(float(xmlbox.find('xmax').text)),
    
              int(float(xmlbox.find('ymin').text)),
    
              int(float(xmlbox.find('ymax').text)))
    
         
    
         # 将坐标转换为YOLO格式
    
         x_center = (b[0] + b[1]) / 2.0
    
         y_center = (b[2] + b[3]) / 2.0
    
         width = b[1] - b[0]
    
         height = b[3] - b[2]
    
         
    
         x_center /= img_width
    
         y_center /= img_height
    
         width /= img_width
    
         height /= img_height
    
         
    
         txt_file.write(f"{cls_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
    
  
    
 # 使用方法
    
 img_width, img_height = 1024, 1024  # 图像尺寸
    
 convert_voc_to_yolo('path/to/image_0001.xml', 'path/to/image_0001.txt', img_width, img_height)
    
    
    
    
    AI写代码

 cd ultralytics

    
 python train.py --img 1024 --batch 16 --epochs 100 --data ../DIOR_dataset/data.yaml --weights yolov8m.pt --cache
    
    
    
    
    AI写代码

 import torch

    
 from ultralytics import YOLO  # 使用YOLOv8的API
    
  
    
 def main():
    
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
     data_yaml = '../DIOR_dataset/data.yaml'
    
     
    
     # 加载预训练模型
    
     model = YOLO('yolov8m.pt').to(device)  # 使用YOLOv8m作为基础模型
    
     
    
     # 设置训练参数
    
     epochs = 100
    
     batch_size = 16
    
     img_size = 1024
    
     
    
     # 开始训练
    
     model.train(data=data_yaml, epochs=epochs, batch=batch_size, imgsz=img_size, device=device)
    
     
    
     # 保存模型
    
     model.save('trained_model.pt')
    
     print('Training complete.')
    
  
    
 if __name__ == '__main__':
    
     main()
    
    
    
    
    AI写代码

    python detect.py --weights runs/train/exp/weights/best.pt --img 1024 --conf 0.4 --source ../DIOR_dataset/valid/images/
    
    AI写代码

    python detect.py --weights best.pt --img 1024 --conf 0.4 --source ../DIOR_dataset/valid/images/
    
    AI写代码