yolov8环境搭建+训练自己的数据集

新建conda环境

conda create -n env_yolov8_bicycle python=3.10

conda activate env_yolov8_bicycle

安装pytorch和ultralytics

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

pip install ultralytics

yolov8源码下载下载

https://github.com/ultralytics/ultralytics

准备数据集

首先在根目录下，新建datasets文件夹
这里用的是随便下载的数据集（绝缘子缺陷检测），将下载的insulator数据集文件夹复制到datasets文件夹中。但是这个数据集是voc格式，需要将其转换为yolo格式

https://aistudio.baidu.com/datasetdetail/122549

voc转yolo

在datasets文件夹下，新建xml2txt.py。注意，需要修改imgpath 、xmlpath、txtpath。

# xml2txt.py
import xml.etree.ElementTree as ET
import os, cv2
import numpy as np
from os import listdir
from os.path import join

classes = []

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def convert_annotation(xmlpath, xmlname):
    with open(xmlpath, "r", encoding='utf-8') as in_file:
        txtname = xmlname[:-4] + '.txt'
        txtfile = os.path.join(txtpath, txtname)
        tree = ET.parse(in_file)
        root = tree.getroot()
        filename = root.find('filename')
        img = cv2.imdecode(np.fromfile('{}/{}.{}'.format(imgpath, xmlname[:-4], postfix), np.uint8), cv2.IMREAD_COLOR)
        h, w = img.shape[:2]
        res = []
        for obj in root.iter('object'):
            cls = obj.find('name').text
            if cls not in classes:
                classes.append(cls)
            cls_id = classes.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((w, h), b)
            res.append(str(cls_id) + " " + " ".join([str(a) for a in bb]))
        if len(res) != 0:
            with open(txtfile, 'w+') as f:
                f.write('\n'.join(res))


if __name__ == "__main__":
    postfix = 'jpg'
    imgpath = 'insulator/JPEGImages' #图片路径，需修改 
    xmlpath = 'insulator/Annotations' #voc标注文件路径，，需修改 
    txtpath = 'insulator/txt' #yolo标注文件路径，需修改 
    
    if not os.path.exists(txtpath):
        os.makedirs(txtpath, exist_ok=True)
    
    list = os.listdir(xmlpath)
    error_file_list = []
    for i in range(0, len(list)):
        try:
            path = os.path.join(xmlpath, list[i])
            if ('.xml' in path) or ('.XML' in path):
                convert_annotation(path, list[i])
                print(f'file {list[i]} convert success.')
            else:
                print(f'file {list[i]} is not xml format.')
        except Exception as e:
            print(f'file {list[i]} convert error.')
            print(f'error message:\n{e}')
            error_file_list.append(list[i])
    print(f'this file convert failure\n{error_file_list}')
    print(f'Dataset Classes:{classes}')

划分train、test、vaild

将数据集从voc转换为yolo格式化后，按照0.7 0.15 0.15的比例，进行分割。在datasets文件夹下，新建xml2txt.py。

# split.py
import os
import random
import shutil


def copy_files(src_dir, dst_dir, filenames, extension):
    os.makedirs(dst_dir, exist_ok=True)
    missing_files = 0
    for filename in filenames:
        src_path = os.path.join(src_dir, filename + extension)
        dst_path = os.path.join(dst_dir, filename + extension)

        # Check if the file exists before copying
        if os.path.exists(src_path):
            shutil.copy(src_path, dst_path)
        else:
            print(f"Warning: File not found for {filename}")
            missing_files += 1

    return missing_files


def split_and_copy_dataset(image_dir, label_dir, output_dir, train_ratio=0.7, valid_ratio=0.15, test_ratio=0.15):
    # 获取所有图像文件的文件名（不包括文件扩展名）
    image_filenames = [os.path.splitext(f)[0] for f in os.listdir(image_dir)]

    # 随机打乱文件名列表
    random.shuffle(image_filenames)

    # 计算训练集、验证集和测试集的数量
    total_count = len(image_filenames)
    train_count = int(total_count * train_ratio)
    valid_count = int(total_count * valid_ratio)
    test_count = total_count - train_count - valid_count

    # 定义输出文件夹路径
    train_image_dir = os.path.join(output_dir, 'train', 'images')
    train_label_dir = os.path.join(output_dir, 'train', 'labels')
    valid_image_dir = os.path.join(output_dir, 'valid', 'images')
    valid_label_dir = os.path.join(output_dir, 'valid', 'labels')
    test_image_dir = os.path.join(output_dir, 'test', 'images')
    test_label_dir = os.path.join(output_dir, 'test', 'labels')

    # 复制图像和标签文件到对应的文件夹
    train_missing_files = copy_files(image_dir, train_image_dir, image_filenames[:train_count], '.jpg')
    train_missing_files += copy_files(label_dir, train_label_dir, image_filenames[:train_count], '.txt')

    valid_missing_files = copy_files(image_dir, valid_image_dir, image_filenames[train_count:train_count + valid_count],
                                     '.jpg')
    valid_missing_files += copy_files(label_dir, valid_label_dir,
                                      image_filenames[train_count:train_count + valid_count], '.txt')

    test_missing_files = copy_files(image_dir, test_image_dir, image_filenames[train_count + valid_count:], '.jpg')
    test_missing_files += copy_files(label_dir, test_label_dir, image_filenames[train_count + valid_count:], '.txt')

    # Print the count of each dataset
    print(f"Train dataset count: {train_count}, Missing files: {train_missing_files}")
    print(f"Validation dataset count: {valid_count}, Missing files: {valid_missing_files}")
    print(f"Test dataset count: {test_count}, Missing files: {test_missing_files}")


# 使用例子
image_dir = 'bicycle_dataset/images'
label_dir = 'bicycle_dataset/labels'
output_dir = './bicycle_dataset_split'

split_and_copy_dataset(image_dir, label_dir, output_dir)

新建yaml文件

在datasets文件夹中，创建 main.yaml 文件。注意，这里要修改path路径

path: C:\AI\demo_yolov8_1\datasets\insulator_split
train: train/images
val: valid/images
test: test/images
# number of classes
nc: 1

# class names
names: ['insulator']

配置Yolov8

切换至ultralytics/cfg/models/v8 目录，修改yolov8.yaml文件，将nc改为1（这里主要修改的是类别数nc，数据集有几种类型，就填几）。

模型训练

下载预训练模型yolo8s.pt，放在datasets下
在Yolov8的GitHub开源库中下载对应版本的模型。我选的是yolo8s.pt。

https://pan.baidu.com/s/1tsOxijvq9auwYw0v5jrXWA?pwd=8888

9. 运行

yolo task=detect mode=train model=datasets/yolov8s.pt epochs=100 batch=10 data=datasets/main.yaml