一、数据标注

1、下载标注工具X-Anylabeling

github地址:https://github.com/CVHub520/X-AnyLabeling

直接下载源码,根据Readme.md配置环境即可

下载并配置环境后

windows cmd打开终端
# 进入文件夹目录
cd D:\AI\X-AnyLabeling-main

python anylabeling/app.py

直接按p打开绘制旋转矩形框即可,zxcv可以调整角度

注意标注完后导出的格式为DOTA 

导出结果如下图所示

二、yolov8源码下载

yolov8源码:GitHub - ultralytics/ultralytics: NEW - YOLOv8 🚀 in PyTorch > ONNX > OpenVINO > CoreML > TFLite

模型:OBB - Ultralytics YOLO Docs

三、数据预处理

如上图所示,这是最终处理之后的目录结构

1、划分数据集

import random
import os
import shutil

def split_dataset(source_images_dir, source_labels_dir, train_images_dir, train_labels_dir, val_images_dir, val_labels_dir, val_ratio=0.2):
    # 创建目标文件夹
    os.makedirs(train_images_dir, exist_ok=True)
    os.makedirs(train_labels_dir, exist_ok=True)
    os.makedirs(val_images_dir, exist_ok=True)
    os.makedirs(val_labels_dir, exist_ok=True)

    # 获取图像文件列表
    image_files = os.listdir(source_images_dir)

    # 计算验证集的大小
    val_size = int(len(image_files) * val_ratio)

    # 随机选择验证集的图像文件
    val_images = set(random.sample(image_files, val_size))

    # 复制图像和标签文件到相应的目标文件夹
    for image_file in image_files:
        image_path = os.path.join(source_images_dir, image_file)
        label_file = image_file.replace('.jpg', '.txt')  # 假设标签文件与图像文件同名,不同后缀
        label_path = os.path.join(source_labels_dir, label_file)

        if image_file in val_images:
            shutil.copy(image_path, os.path.join(val_images_dir, image_file))
            shutil.copy(label_path, os.path.join(val_labels_dir, label_file))
        else:
            shutil.copy(image_path, os.path.join(train_images_dir, image_file))
            shutil.copy(label_path, os.path.join(train_labels_dir, label_file))

# 调用函数来划分数据集
split_dataset(r'D:\1_ggjx\work\trough\obb\JPEGImages', r'D:\1_ggjx\work\trough\obb\org\labelTxt', r'dataset/images\train', r'dataset\labels\train_original', r'dataset\images\val', r'dataset\labels\val_original', val_ratio=0.2)

修改为自己的路径即可,其中C:\Users\Yaogun-Ma\Desktop\trough\obb\JPEGImages和C:\Users\Yaogun-Ma\Desktop\trough\obb\labelTxt是我标注的数据集路径,其它的都是需要存放划分好的数据,建议跟我这个保持一样,之后自己再做调整都行

2、DOTA数据转为yolov8-obb所训练的格式

from ultralytics.data.converter import convert_dota_to_yolo_obb




convert_dota_to_yolo_obb(r'D:\AI\obb\ultralytics-main\dataset')

路径改为自己刚刚创建的路径,然后进入converter.py修改类别,把原来的注释掉改为自己的类别,另外就是下方有个图片格式要自己确认,里面默认用的是png,如果是jpg的话就要修改一下。

然后开始转换格式,运行结果如下

最后就会在我们之前划分数据的路径中的labels文件夹创建了train和val这两个文件了

四、训练

1、修改yaml文件

在路径ultralytics-main/ultralytics/cfg/datasets下新建一个mydata.yaml文件

# Ultralytics YOLO 🚀, AGPL-3.0 license
# DOTA 1.0 dataset https://captain-whu.github.io/DOTA/index.html for object detection in aerial images by Wuhan University
# Documentation: https://docs.ultralytics.com/datasets/obb/dota-v2/
# Example usage: yolo train model=yolov8n-obb.pt data=DOTAv1.yaml
# parent
# ├── ultralytics
# └── datasets
#     └── dota1  ← downloads here (2GB)

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: D:\AI\obb\ultralytics-main\dataset # dataset root dir
train: images/train # train images (relative to 'path') 1411 images
val: images/val # val images (relative to 'path') 458 images
test: images/test # test images (optional) 937 images

# Classes for DOTA 1.0
names:
  0: trough

把代码复制进去,修改自己的路径和类别即可

修改自己的类被数量ultralytics/cfg/models/v8/yolov8-obb.yaml

2、开始train

from ultralytics import YOLO


def main():
    model = YOLO(r'D:\AI\obb\ultralytics-main\ultralytics\cfg\models\v8\yolov8-obb.yaml').load(r'D:\AI\obb\ultralytics-main\yolov8n-obb.pt')  # build from YAML and transfer weights
    model.train(data='./ultralytics/cfg/datasets/mydata.yaml', epochs=100, imgsz=640, batch=4, workers=0)


if __name__ == '__main__':
    main()

修改自己的预训练权重路径和其它路径,再运行上面的代码即可

3、测试

from ultralytics import YOLO

# Load a model
# model = YOLO("yolov8n-obb.pt")  # load an official model
model = YOLO(r"D:\AI\obb\ultralytics-main\runs\obb\train3\weights\best.pt")  # load a custom model

# Predict with the model
results = model(r"D:\AI\obb\ultralytics-main\dataset\images\train", save=True)  # predict on an image

修改自己训练后的权重路径、测试数据路即可

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐