训练自己的深度学习模型（适用于YOLOv5、YOLOv6、YOLOv7、YOLOv8、YOLOv9）

训练自己的深度学习模型

数据集划分代码：https://github.com/Jie-Huangi/dataset_labelImg_To_yolo，欢迎标星🔥🔥🔥。当前仓库的代码还不够完整，后期会添加数据增强相关代码。

1、收集数据

数据来源可以是，网络上复制粘贴，也可以自己写代码爬取，其次就是手机拍摄，摄像头拍摄，还有网上下载别人整理好的数据集。

具体的数据怎么来，根据自己的情况安排就可以。

我这里收集的奶龙和小七图片是自己在B站上，截屏获取的。

唯一需要注意的是，图片格式最好时 png、jpg等常见的格式，不然后期算法识别不到对应的数据，就会报错。

2、标注数据，labelme标注软件。

https://github.com/HumanSignal/labelImg

3、制作数据集

使用labelImg标注软件，获取物体位置信息后，得到image文件和txt文件。为了适应yolo算法读取数据，需要调整images和labels文件的存放格式。数据存放格式如下面代码中展示的形式。要得到这种格式，可以手动复制粘贴，如果遇到成千上百份文件，更方便的还是使用代码，自动生成。调整完成后，即可完成数据集的制作。

# ├── datasets
#     └── mydata  
#         └── train
#             └── images
#             └── labels
#         └── val
#             └── images
#             └── labels
#         └── test

import argparse
import glob
from pathlib import Path
import random
import shutil
import os
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor
NUM_THREADS = min(8, max(1, os.cpu_count() - 1))
def run(func, this_iter, desc="Processing"):
    with ThreadPoolExecutor(max_workers=NUM_THREADS, thread_name_prefix='MyThread') as executor:
        results = list(
            tqdm(executor.map(func, this_iter), total=len(this_iter), desc=desc)
        )
    return results
def split_dataset_into_train_val_test(
    dataset_dir,
    save_dir,
    train_ratio=0.7,
    val_ratio=0.2,
    test_ratio=0.1,
    im_suffix='jpg'
):
    if isinstance(dataset_dir, str):
        dataset_dir = Path(dataset_dir)
    image_files = []
    for suffix in im_suffix:
        image_files += glob.glob(str(dataset_dir / 'images' / f"*.{suffix}"))
    total_images = len(image_files)
    random.shuffle(image_files)
    train_split = int(total_images * train_ratio)
    val_split = int(total_images * val_ratio)
    # test_split = int(total_images * test_ratio)
    if train_ratio + val_ratio == 1:
        train_images = image_files[:train_split]
        val_images = image_files[train_split:]
        test_images = []
    else:
        train_images = image_files[:train_split]
        val_images = image_
files[train_split : train_split + val_split]
        test_images = image_files[train_split + val_split :]
    print('*'*25)
    print(
        "",
        f"Total images: {total_images}\n",
        f"Train images: {len(train_images)}\n",
        f"Val images: {len(val_images)}\n",
        f"Test images: {len(test_images)}"
    )
    print('*'*25)
    split_paths = [("train", train_images), ("val", val_images), ("test", test_images)]
    for split_name, images in split_paths:
        split_dir = Path(save_dir) / split_name
        for dir_name in ['images', 'labels']:
            if not (split_dir / dir_name).exists():
                (split_dir / dir_name).mkdir(exist_ok=True, parents=True)
        args_list = [(image, dataset_dir, split_dir) for image in images]
        run(process_image, args_list, desc=f"Creating {split_name} dataset")
        print(f"Created {split_name} dataset with {len(images)} images.")
def process_image(args):
    image_file, dataset_dir, split_dir = args
    annotation_file = dataset_dir / 'labels' / f"{Path(image_file).stem}.txt"
    assert annotation_file.exists(), f'{annotation_file} 不存在!'
    if not has_objects(annotation_file):
        return
    shutil.copy(image_file, split_dir / "images" / Path(image_file).name)
    shutil.copy(annotation_
file, split_dir / "labels" / annotation_file.name)
def has_objects(annotation_path):
    with open(annotation_path, "r") as f:
        lines = f.readlines()
    return len(lines) > 0

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', default='./data')  # 数据集Images路径
    parser.add_argument('--save', default='./mydata')  # 保存路径
    parser.add_argument('--images_suffix', default=['jpg', 'png', 'jpeg'], help='images suffix')  # 图片后缀名
    opt = parser.parse_args()
    split_dataset_into_train_val_test(
        dataset_dir=opt.data,
        save_dir=opt.save,
        train_ratio=0.8,
        val_ratio=0.2,
        im_suffix=opt.images_suffix
    )

4、配置训练文件mycoco.yaml

配置文件中的path为数据集的根目录。train位置可以设置为images的文件夹路径，也可以是每张图片的txt路径。同理，val也是。

# ├── datasets
#     └── mydata
#         └── train
#             └── images
#             └── labels
#         └── val
#             └── images
#             └── labels
#         └── test

# 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: /home/bruce/Documents/Bilibili_Record/dataset/mydata # dataset root dir
train: train/images # train images (relative to 'path') 128 images
val: val/images # val images (relative to 'path') 128 images
test: # test images (optional)

# Classes
names:
  0: nailong
  1: xiaoqi

5、设置训练文件train.py

训练文件，train.py中，只需要设置好预训练权重文件，数据集文件位置，还有迭代次数，以及batch-size的大小。这些参数，刚开始根据个人电脑性能进行设置。主要目的是学会如何训练自己的数据集。在写论文，做实验的时候，就需要参考论文中的参数进行设置。

6、detect 检测结果

训练完成后，会得到一个权重文件，在runs/train/exp/weights路径下，找到最新pt文件，复制该路径到detect.py的参数设置模块。主要需要设置的参数，就是weights、source、data。具体设置下图中所示。

设置完成后，点击运行该代码，即可以预测source中的图片。

7、val 评估结果

当我们训练好一个模型过后，需要评估模型的性能，比如查看平均检测精度（mAP），召回率（R），精确率（P）等参数，就可以使用这个val.py文件。

需要设置的参数主要涉及三个位置，数据集的位置，权重（weights）文件的位置，batch-size的大小。具体设置如下图所示。