yolov8训练自己的目标检测数据集详细介绍版

pip install ultralytics wandb

# data.yaml
path: ../datasets/MyDataSet  # 数据集根路径
train: ./train  # 训练集路径
val: ./val  # 验证集路径
nc: 2  # 类别数量
names: ['class1', 'class2']  # 类别名称

from ultralytics import YOLO

# 加载预训练模型
model = YOLO('yolov8n.pt')

# 定义数据集配置文件路径
data_config = 'path/to/your/data.yaml'  # 替换为你的data.yaml文件的实际路径

# 设置训练参数
epochs = 100
imgsz = 640
batch_size = 16
workers = 8
device = 0  # 使用GPU编号，如果使用CPU则设为None或'cpu'

# 开始训练
results = model.train(
    data=data_config,
    epochs=epochs,
    imgsz=imgsz,
    batch=batch_size,
    workers=workers,
    device=device,
    name='yolov8_custom_train',  # 模型保存名称
    save=True,  # 是否保存模型
    exist_ok=True,  # 如果存在相同命名的训练结果是否覆盖
    cache=False  # 是否缓存数据以加速训练
)

print("训练完成:", results.save_dir)

import os
import shutil
import xml.etree.ElementTree as ET

# VOC格式数据集路径
voc_data_path = r'E:\DataSet\VOC2012'
voc_annotations_path = os.path.join(voc_data_path, 'Annotations')
voc_images_path = os.path.join(voc_data_path, 'JPEGImages')

# YOLO格式数据集保存路径
yolo_data_path = 'E:\\DataSet\\helmet-YOLO'
yolo_images_path = os.path.join(yolo_data_path, 'images')
yolo_labels_path = os.path.join(yolo_data_path, 'labels')

# 创建YOLO格式数据集目录
os.makedirs(yolo_images_path, exist_ok=True)
os.makedirs(yolo_labels_path, exist_ok=True)

# 类别映射 (可以根据自己的数据集进行调整)
class_mapping = {
    'dog': 0,
    'cat': 1,
    # 添加更多类别...
}

def convert_voc_to_yolo(voc_annotation_file, yolo_label_file):
    tree = ET.parse(voc_annotation_file)
    root = tree.getroot()

    size = root.find('size')
    width = float(size.find('width').text)
    height = float(size.find('height').text)

    with open(yolo_label_file, 'w') as f:
        for obj in root.findall('object'):
            cls = obj.find('name').text
            if cls not in class_mapping:
                continue
            cls_id = class_mapping[cls]
            xmlbox = obj.find('bndbox')
            xmin = float(xmlbox.find('xmin').text)
            ymin = float(xmlbox.find('ymin').text)
            xmax = float(xmlbox.find('xmax').text)
            ymax = float(xmlbox.find('ymax').text)

            x_center = (xmin + xmax) / 2.0 / width
            y_center = (ymin + ymax) / 2.0 / height
            w = (xmax - xmin) / width
            h = (ymax - ymin) / height

            f.write(f"{cls_id} {x_center} {y_center} {w} {h}\n")

# 遍历VOC数据集的Annotations目录，进行转换
for voc_annotation in os.listdir(voc_annotations_path):
    if voc_annotation.endswith('.xml'):
        voc_annotation_file = os.path.join(voc_annotations_path, voc_annotation)
        image_id = os.path.splitext(voc_annotation)[0]
        voc_image_file = os.path.join(voc_images_path, f"{image_id}.jpg")
        yolo_label_file = os.path.join(yolo_labels_path, f"{image_id}.txt")
        yolo_image_file = os.path.join(yolo_images_path, f"{image_id}.jpg")

        convert_voc_to_yolo(voc_annotation_file, yolo_label_file)
        if os.path.exists(voc_image_file):
            shutil.copy(voc_image_file, yolo_image_file)

print("转换完成！")

import os
import shutil
import random

def make_yolo_dataset(images_folder, labels_folder, output_folder, train_ratio=0.8):
    # 创建目标文件夹
    images_train_folder = os.path.join(output_folder, 'images/train')
    images_val_folder = os.path.join(output_folder, 'images/val')
    labels_train_folder = os.path.join(output_folder, 'labels/train')
    labels_val_folder = os.path.join(output_folder, 'labels/val')

    os.makedirs(images_train_folder, exist_ok=True)
    os.makedirs(images_val_folder, exist_ok=True)
    os.makedirs(labels_train_folder, exist_ok=True)
    os.makedirs(labels_val_folder, exist_ok=True)

    # 获取图片和标签的文件名（不包含扩展名）
    image_files = [f for f in os.listdir(images_folder) if f.endswith('.jpg')]
    label_files = [f for f in os.listdir(labels_folder) if f.endswith('.txt')]
    image_base_names = set(os.path.splitext(f)[0] for f in image_files)
    label_base_names = set(os.path.splitext(f)[0] for f in label_files)

    # 找出图片和标签都存在的文件名
    matched_files = list(image_base_names & label_base_names)

    # 打乱顺序并划分为训练集和验证集
    random.shuffle(matched_files)
    split_idx = int(len(matched_files) * train_ratio)
    train_files = matched_files[:split_idx]
    val_files = matched_files[split_idx:]

    # 移动文件到对应文件夹
    for base_name in train_files:
        img_src = os.path.join(images_folder, f"{base_name}.jpg")
        lbl_src = os.path.join(labels_folder, f"{base_name}.txt")

        img_dst = os.path.join(images_train_folder, f"{base_name}.jpg")
        lbl_dst = os.path.join(labels_train_folder, f"{base_name}.txt")

        shutil.copyfile(img_src, img_dst)
        shutil.copyfile(lbl_src, lbl_dst)

    for base_name in val_files:
        img_src = os.path.join(images_folder, f"{base_name}.jpg")
        lbl_src = os.path.join(labels_folder, f"{base_name}.txt")

        img_dst = os.path.join(images_val_folder, f"{base_name}.jpg")
        lbl_dst = os.path.join(labels_val_folder, f"{base_name}.txt")

        shutil.copyfile(img_src, img_dst)
        shutil.copyfile(lbl_src, lbl_dst)

    print("数据集划分完成！")

# 使用示例
images_folder = 'path/to/your/images_folder'  # 原始图片文件夹路径
labels_folder = 'path/to/your/labels_folder'  # 原始标签文件夹路径
output_folder = 'path/to/your/output_folder'  # 存放结果数据集的文件夹路径
make_yolo_dataset(images_folder, labels_folder, output_folder)

train: E:/VOC2012/train/images  # train images (relative to 'path') 128 images
val: E:/VOC2012/val/images  # val images (relative to 'path') 128 images
test: E:/VOC2012/test/images
 
nc: 2
 
# Classes
names: ['dog','cat']

from ultralytics import YOLO
#多线程添加代码
if __name__ == '__main__': 
    model = YOLO('yolov8n.yaml').load('yolov8n.pt')
    # # 训练模型
    results = model.train(data='data.yaml', epochs=100, imgsz=640, device=0, workers=2, batch=16, optimizer='SGD', amp=False)

from ultralytics import YOLO
# 加载训练好的模型，改为自己的路径
model = YOLO('runs/detect/train/weights/best.pt')
# 修改为自己的图像或者文件夹的路径
source = 'test.jpg' #修改为自己的图片路径及文件名
# 运行推理，并附加参数
model.predict(source, save=True)