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

深度学习模型训练指南：从数据收集到模型评估

原创

已于 2024-05-16 09:37:49 修改 · 1.2k 阅读

23 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能 #YOLO #python #目标检测 #计算机视觉

于 2024-05-14 10:15:14 首次发布

训练自己的深度学习模型

数据集划分代码：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