YOLOV8训练自定义数据集

一、 仓库地址

YOLOV8

二、数据预处理

2.1 将标注数据按照类别划分到不同的文件目录

# 单独筛选某一类缺陷
import numpy as np
import os
#shutil:操作多文件或者文件集合（复制、移除等）
import shutil

#------------------------------------------------------------------#
#作用：
#   按照类别进行数据集的存储
# 输入：
# 1）存储txt文件的地址
# 2）输出文件目录
#------------------------------------------------------------------#
def move_file(src, out_path):
    data = []
    # 设置文件对象并读取每一行文件，测试集路径
    for line in open(src, "r"):
        data.append(line)
    for a in data:  # 读取每一行
        #选中标注的类别
        cls = a[0]
        #print(cls)
        if cls == '11':
            #将 源地址的文件 移动到 目的地址
            shutil.move('.' + src.split('.')[1] + '.txt', out_path)
            shutil.move('.'  + src.split('.')[1] + '.jpg', out_path)
            break

#------------------------------------------------------------------#
#作用：
#   将根目录下的文件地址存到相应的列表中
# 输入：
# 1）待分类图像的根目录
# 2）存储文件地址的列表
# 3）存储文件夹地址的列表
#------------------------------------------------------------------#
def get_file_path(root_path, file_list, dir_list):

    dir_or_files = os.listdir(root_path)
    for dir_file in dir_or_files:
        # 获取目录或者文件的路径
        dir_file_path = os.path.join(root_path, dir_file)
        # 判断该路径为文件还是路径
        if os.path.isdir(dir_file_path):
            dir_list.append(dir_file_path)
            # 递归获取所有文件和目录的路径
            get_file_path(dir_file_path, file_list, dir_list)
        else:
            file_list.append(dir_file_path);

#------------------------------------------------------------------#
#作用：
#   选中以txt结尾的文件，并按照类别进行划分
# 输入：
# 1）存储文件地址的列表
# 2）输出文件目录
#------------------------------------------------------------------#
def get_file(file_list, out_path):
    for file_name in file_list:
        if file_name.endswith(".txt"):
            #此时file_name 会选中以 txt 结尾的文件
            move_file(file_name, out_path)


if __name__ == "__main__":
    # 指定标注好的数据集路径与输出文件目录
    root_path = r"./datasets/HONOR/MyData/original_Image/"
    out_path = r"./datasets/HONOR/MyData/11/"

    # 用来存放所有的文件及目录路径
    file_list = []
    dir_list = []

    #按照类别进行数据集存储
    get_file_path(root_path, file_list, dir_list)
    get_file(file_list, out_path)

2.2 按照比例划分训练集与验证集

import os
import shutil
import random

'''
目录下随机选择一个文件
Args:
    dir:目录
return:
    filename:随机选择的文件名（绝对路径）
'''
def randomChoiceIn(dir):
    #os.walk 输出 三元组 （根目录、目录列表、文件列表）
    for root,dirs,files in os.walk(dir):
        index=random.randint(0,len(files)-1)
        filename = files[index]
    return root + filename

'''
将选中的文件成对的放到对应文件目录中
Args:
    src : 选中的文件
    out_path_images : 输出图像的文件目录
    out_path_labels : 输出标签的文件目录
'''
def move_file(src, out_path_images, out_path_labels):
    #将 源地址的文件 移动到 目的地址
    shutil.move('.' + src.split('.')[1] + '.txt', out_path_labels)
    shutil.move('.' + src.split('.')[1] + '.jpg', out_path_images)

'''
将根目录下的文件地址存到相应的列表中
Args:
    root_path : 待分类图像的根目录
    file_list : 存储文件地址的列表
    dir_list : 存储文件夹地址的列表
'''
def get_file_path(root_path, file_list, dir_list):

    dir_or_files = os.listdir(root_path)
    for dir_file in dir_or_files:
        # 获取目录或者文件的路径
        dir_file_path = os.path.join(root_path, dir_file)
        # 判断该路径为文件还是路径
        if os.path.isdir(dir_file_path):
            dir_list.append(dir_file_path)
            # 递归获取所有文件和目录的路径
            get_file_path(dir_file_path, file_list, dir_list)
        else:
            file_list.append(dir_file_path);

if __name__ == "__main__":
    # 验证集所占的比例
    valid_percent = 0.15

    #指定输入文件目录
    input_path = r"./datasets/HONOR/MyData/4/"

    # 指定验证集与训练集样本及标签的保存路径
    val_out_path_images = r"./datasets/HONOR/MyData/images/valid0620/"
    val_out_path_labels = r"./datasets/HONOR/MyData/labels/valid0620/"

    train_out_path_images = r"./datasets/HONOR/MyData/images/train0620/"
    train_out_path_labels = r"./datasets/HONOR/MyData/labels/train0620/"

    # 用来存放训练集所有的文件及目录路径
    train_file_list = []
    train_dir_list = []

    # 确定验证集的数量
    dir_or_files = os.listdir(input_path)
    val_num = int(len(dir_or_files) / 2 * valid_percent)

    # 将选中的数据放到验证集对应的样本及标签目录
    for i in range(val_num):
        #随机选中一个文件
        random_file = randomChoiceIn(input_path)

        #将文件与标签分别放到对应的目录
        move_file(random_file, val_out_path_images, val_out_path_labels)

    #获取剩余的训练集图像列表
    get_file_path(input_path, train_file_list, train_dir_list)
    for train_file_name in train_file_list:
        #可以选择文本或者图像数据，避免选中两次
        if train_file_name.endswith(".txt"):
            move_file(train_file_name, train_out_path_images, train_out_path_labels)

2.3 数据集按照既定格式进行整理

三、训练与推理

3.1 连接到远程服务器

git bash here
ssh [email protected]

3.2 使用linux的screen命令打开虚拟窗口，避免因断电、断网而中断训练

#新建虚拟窗口
screen -S name

#查看所有的screen
screen -ls

#恢复screen
screen -r name

#推出
screen -S name -X quit

3.3 进入对应的虚拟环境

#查看对应的虚拟环境
conda info --env

#激活虚拟环境
conda activate name_env

#退出
conda deactivate

3.4 安装ultralytics包

pip install ultralytics #包的路径在home/user目录下

如果安装的Torch和TorchVision是CPU版本的，需要从这里下载对应的GPU版本，通过pip install XXX.whl完成安装。

3.5 开始训练，训练结果在run目录下

yolo task=detect \
mode=train \
model={HOME}/runs/detect/train/weights/best.pt \  //注意 一定是pt的预训练文件、如果是从头训练的yaml文件，难以收敛
data={dataset.location}/data.yaml

3.6 新开一个bash连接到远程服务器上并打开TensorBoard观察模型的训练过程

#默认端口6006，还可以通过--port指定端口
tensorboard --logdir runs/train/exp13

3.7 在本地观察服务器端模型的训练过程

#新开一个bash建立本地与远程的映射
ssh -L 16006:127.0.0.1:6006 [email protected]

#在本地电脑的浏览器输入下边的指令
http://127.0.0.1:16006/

3.8 python单张预测推理

yolo task=detect \
mode=predict \
model={HOME}/runs/detect/train/weights/best.pt \
conf=0.25 \
source={dataset.location}/test/images