Yolov8 obb从零开始训练自己的数据集

目前在打一个比赛，偶然发现了这个yolov8obb效果很好，但是网上的教程有的看不懂！！于是照猫画虎试着去部署了一下。闲话少说，进入正题：

1、部署环境（默认conda）

2、下载项目

git clone github.com/ultralytics/ultralytics

3、进入下载的项目并安装ultralytics

太慢的话就换源-i https://pypi.tuna.tsinghua.edu.cn/simple

4、数据集处理

YOLOv8 OBB支持的数据集格式是 :  Index_number, x, y, x, y, x, y, x, y 就是类别码和四个坐标

而现在我的数据集格式是945.0 999.0 945.0 985.0 985.0 985.0 985.0 999.0 ship 0，就需要对数据集进行转化

！！注意YOLO的数据集格式的坐标是要归一化的！！！（踩坑）

以下是我自己写的函数（代码能力菜别骂别骂）完整代码在最后

4.1数据检查（如果数据集没问题就可以略过）

def examine(label_path, img_path):
    label_list = []
    for file_name in os.listdir(img_path):
        if file_name.endswith('.jpg'): #这里改成自己的图片格式
            label_list.append(file_name.replace('.jpg', '.txt'))

    for file_name in os.listdir(label_path):
        if file_name.endswith('.txt'):
            if file_name in label_list :
                label_list.remove(file_name)
            else :
                print("ERROR: "+file_name+"s image can not be found!")
                return 0
    if len(label_list) > 0:
        print("ERROR: ",label_list, " can not found!")
    else:
        print("The dataset has been examined!")
        return 1

写了一个检查数据集的代码，确保了图片和标签文件的一一对应

4.2数据集转化

def convert_to_yolo_normalize(label_path, target_path, img_path):
    for file_name in os.listdir(label_path):
        if file_name.endswith('.txt'):
            #为了归一化读取对应的图片数据
            image_path = os.path.join(img_path, file_name.replace('.txt', '.jpg'))
            image = Image.open(image_path)
            width, height = image.size
            target_file = os.path.join(target_path, file_name)
            source_file = os.path.join(label_path, file_name)
            with open(source_file, 'r') as source, open(target_file, 'w') as target:
                lines = source.readlines()
                for line in lines:
                    data = line.strip().split()
                    #我这里是一个类，如果是很多类的话就要使用一个映射了
                    #比如line = map(data[-1]).join(data)
                    data = data[:-2]
                    for i in range(0, 8, 2):
                        data[i] = format(float(data[i]) / width , '.4f')
                    for i in range(1, 9, 2):
                        data[i] = format(float(data[i]) / height , '.4f')
                    line = " ".join(data)
                    target.write('0 '+line + '\n')
    print("convert finish")

4.3 数据集划分

由于YOLO对于这个数据集的目录格式有要求，具体要求如下：

那么现在我们的图片和标签各自在一个文件夹当中，需要把他们划分：

def split_dataset(label_path, img_path, dataset_path, ratio_train, ratio_val):
    ratio_test = 1 - ratio_train - ratio_val

    # 获取所有图像文件和标签文件
    all_images = [i for i in os.listdir(img_path) if i.endswith('.jpg')]
    random.shuffle(all_images)
    #开辟数据目录
    images_dir = os.path.join(dataset_path, 'images/')
    labels_dir = os.path.join(dataset_path, 'labels/')
    for dir in [images_dir, labels_dir]:
        os.makedirs(os.path.join(dir, 'train'), exist_ok = True)
        os.makedirs(os.path.join(dir, 'test'), exist_ok = True)
        os.makedirs(os.path.join(dir, 'val'), exist_ok = True)
    # 计算每个集合的大小
    images_counts = len(all_images)
    train_counts = int(images_counts * ratio_train)
    val_counts = int(images_counts * ratio_val)
    #数据集划分
    train_images = all_images[:train_counts]
    val_images = all_images[train_counts:train_counts + val_counts]
    test_images = all_images[train_counts + val_counts:]
    
    for file_name in train_images:
        shutil.copy(os.path.join(img_path, file_name), os.path.join(dataset_path, 'images/train/' ,file_name))
        file_name = file_name.replace('.jpg', '.txt')
        shutil.copy(os.path.join(label_path, file_name), os.path.join(dataset_path, 'labels/train/' ,file_name))
    
    for file_name in val_images:
        shutil.copy(os.path.join(img_path, file_name), os.path.join(dataset_path, 'images/val/' ,file_name))
        file_name = file_name.replace('.jpg', '.txt')
        shutil.copy(os.path.join(label_path, file_name), os.path.join(dataset_path, 'labels/val/' ,file_name))

    for file_name in test_images:
        shutil.copy(os.path.join(img_path, file_name), os.path.join(dataset_path, 'images/test/' ,file_name))
        file_name = file_name.replace('.jpg', '.txt')
        shutil.copy(os.path.join(label_path, file_name), os.path.join(dataset_path, 'labels/test/' ,file_name))
    print("dataset has been splitted!")

5.构建数据集yaml

这个文件是用于告诉模型数据集的相关情况，在工作目录下创建一个dota.yaml文件

path: /root/YOLOv8obb/datasar/dataset # 数据集根目录
train: images/train
val: images/val
test: images/test
names:
  0: ship

在这里我只有一个类，因此只有一个0，如果有多个的话要都加进去

6、下载权重文件

wget https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8x-obb.pt

权重文件可以自己在官网查看下载哦，直接搜索YOLOv8OBB即可

7、运行

yolo obb train data=dota.yaml model=yolov8x-obb.pt epochs=100 imgsz=640

第一次启动可能会下载点东西，现在下载的就是字体

现在就是在正常训练

完整代码(构建数据集）

import os
import random
from PIL import Image
import shutil

label_path = "/root/YOLOv8obb/datasar/data_fix/labels/"
target_path = "/root/YOLOv8obb/datasar/data_fix/yololabels/"
img_path = "/root/YOLOv8obb/datasar/data_fix/images/"
dataset_path = "/root/YOLOv8obb/datasar/dataset/"

def examine(label_path, img_path):
    label_list = []
    for file_name in os.listdir(img_path):
        if file_name.endswith('.jpg'): #这里改成自己的图片格式
            label_list.append(file_name.replace('.jpg', '.txt'))

    for file_name in os.listdir(label_path):
        if file_name.endswith('.txt'):
            if file_name in label_list :
                label_list.remove(file_name)
            else :
                print("ERROR: "+file_name+"s image can not be found!")
                return 0
    if len(label_list) > 0:
        print("ERROR: ",label_list, " can not found!")
    else:
        print("The dataset has been examined!")
        return 1

def convert_to_yolo(label_path, target_path):
    for file_name in os.listdir(label_path):
        if file_name.endswith('.txt'):
            target_file = os.path.join(target_path, file_name)
            source_file = os.path.join(label_path, file_name)
            with open(source_file, 'r') as source, open(target_file, 'w') as target:
                lines = source.readlines()
                for line in lines:
                    data = line.strip().split()
                    #我这里是一个类，如果是很多类的话就要使用一个映射了
                    #比如line = map(data[-1]).join(data)
                    data = data[:-2]
                    line = " ".join(data)
                    target.write('0 '+line + '\n')
    print("convert finish")

def convert_to_yolo_normalize(label_path, target_path, img_path):
    for file_name in os.listdir(label_path):
        if file_name.endswith('.txt'):
            #为了归一化读取对应的图片数据
            image_path = os.path.join(img_path, file_name.replace('.txt', '.jpg'))
            image = Image.open(image_path)
            width, height = image.size
            target_file = os.path.join(target_path, file_name)
            source_file = os.path.join(label_path, file_name)
            with open(source_file, 'r') as source, open(target_file, 'w') as target:
                lines = source.readlines()
                for line in lines:
                    data = line.strip().split()
                    #我这里是一个类，如果是很多类的话就要使用一个映射了
                    #比如line = map(data[-1]).join(data)
                    data = data[:-2]
                    for i in range(0, 8, 2):
                        data[i] = format(float(data[i]) / width , '.4f')
                    for i in range(1, 9, 2):
                        data[i] = format(float(data[i]) / height , '.4f')
                    line = " ".join(data)
                    target.write('0 '+line + '\n')
    print("convert finish")

def split_dataset(target_path, img_path, dataset_path, ratio_train, ratio_val):
    ratio_test = 1 - ratio_train - ratio_val

    # 获取所有图像文件和标签文件
    all_images = [i for i in os.listdir(img_path) if i.endswith('.jpg')]
    random.shuffle(all_images)
    #开辟数据目录
    images_dir = os.path.join(dataset_path, 'images/')
    labels_dir = os.path.join(dataset_path, 'labels/')
    for dir in [images_dir, labels_dir]:
        os.makedirs(os.path.join(dir, 'train'), exist_ok = True)
        os.makedirs(os.path.join(dir, 'test'), exist_ok = True)
        os.makedirs(os.path.join(dir, 'val'), exist_ok = True)
    # 计算每个集合的大小
    images_counts = len(all_images)
    train_counts = int(images_counts * ratio_train)
    val_counts = int(images_counts * ratio_val)
    #数据集划分
    train_images = all_images[:train_counts]
    val_images = all_images[train_counts:train_counts + val_counts]
    test_images = all_images[train_counts + val_counts:]
    
    for file_name in train_images:
        shutil.copy(os.path.join(img_path, file_name), os.path.join(dataset_path, 'images/train/' ,file_name))
        file_name = file_name.replace('.jpg', '.txt')
        shutil.copy(os.path.join(target_path, file_name), os.path.join(dataset_path, 'labels/train/' ,file_name))
    
    for file_name in val_images:
        shutil.copy(os.path.join(img_path, file_name), os.path.join(dataset_path, 'images/val/' ,file_name))
        file_name = file_name.replace('.jpg', '.txt')
        shutil.copy(os.path.join(target_path, file_name), os.path.join(dataset_path, 'labels/val/' ,file_name))

    for file_name in test_images:
        shutil.copy(os.path.join(img_path, file_name), os.path.join(dataset_path, 'images/test/' ,file_name))
        file_name = file_name.replace('.jpg', '.txt')
        shutil.copy(os.path.join(target_path, file_name), os.path.join(dataset_path, 'labels/test/' ,file_name))
    print("dataset has been splitted!")

convert_to_yolo_normalize(label_path, target_path, img_path)
split_dataset(target_path, img_path, dataset_path, 0.8, 0.1)

至此一个简单的部署步骤就完成啦！

我也是cv的新手希望各位大佬不吝赐教，也欢迎各位一起讨论