从COCO等公开数据集自制数据集，实现YOLOv5等目标检测模型训练

初衷

每次训练自制数据集时总要写个小脚本，长此以往不如系统地整理一个吧。该脚本可从公开数据集抽取自己想要目标训练，例如想进行人的检测，可从COCO等各种带txt格式数据集中，把人单独抽出来，拼接成一个更大的人的数据集，实现单目标检测的训练。该方法同样支持自由组合多种标签训练。

COCO数据集txt格式的label下载

链接: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco2017labels.zip

Python脚本

如果tqdm已安装，会以进度条形式输出，未安装则显示百分比
依据注释更改对应路径，组织文件夹结构即可使用

import sys
import os
try:
    from tqdm import tqdm
    module_found = True
except:
    module_found = False

def label_txt(dataset):
    if dataset == "train":
        folder_official = "D:/Data/2D/database/yolov5/coco/labels_default/" # change the path according to your original dataset such as coco2017
        folder_custom = "D:/Data/2D/database/yolov5/coco/labels/" # change the path according to the dataset you want to generate
        folder_type = "train2017" #
    if dataset == "val":
        folder_official = "D:/Data/2D/database/yolov5/coco/labels_default/" # change the path according to your original dataset such as coco2017
        folder_custom = "D:/Data/2D/database/yolov5/coco/labels/" # change the path according to the dataset you want to generate
        folder_type = "val2017" #

    files = list(os.listdir(folder_official+folder_type))
    print(files)
    cnt_obj = 0
    all_files = len(files)
    cnt_files = 0

    if module_found:
        process_bar = tqdm(total=all_files)
        process_bar.set_description('Processing:')
    for file in files:
        custom_label = [] 
        open_txt = False
        if module_found:
            process_bar.update(1)
        else:
            cnt_files += 1
            progress = str("%.2f" % (cnt_files/all_files*100))
            sys.stdout.write('\r'+ f"Progress:{progress}%")

        with open(os.path.join(folder_official+folder_type, file), 'r') as f:
            strs = [x.split() for x in f.read().strip().splitlines()]
            for single_line in strs:
                if single_line[0] == '0': # change "0" to the label you want; modifying the codes to add labels accordingly; "0" is the label of "person" in coco2017
                    custom_label.append(single_line)
                    cnt_obj += 1
                    open_txt = True
            f.close()
        if open_txt == True:
            with open(os.path.join(folder_custom+folder_type, file), 'w') as fp:
                for line in custom_label:
                    newline = " ".join(line) + "\n"
                    fp.writelines(newline)
                fp.close()
            
    print(f'\nnumber of target:{cnt_obj}')

label_txt("train")
label_txt("val")

def whole_txt(dataset):
    if dataset == "train":
        folder = "train2017"
        files = list(os.listdir(folder))
        with open('../train2017_person.txt', 'w') as ftrain:
            for file in files:
                line = './images/train2017/' + file.replace("txt", "jpg") + '\n'
                ftrain.writelines(line)
            ftrain.close()
    
    if dataset == "val":
        folder = "val2017"
        files = list(os.listdir(folder))
        with open('../val2017_person.txt', 'w') as fval:
            for file in files:
                line = './images/val2017/' + file.replace("txt", "jpg") + '\n'
                fval.writelines(line)
            fval.close()

whole_txt("train")
whole_txt("val")

import shutil
def create_img_folder():
    img_folder = "D:/Data/2D/database/yolov5/coco/images/train2017/"
    new_img_folder = "D:/Data/2D/database/yolov5/coco_person/images/train2017/"
    folder = "train2017"
    files = list(os.listdir(folder))
    all_files = len(files)
    if not os.path.exists(new_img_folder):
        os.makedirs(new_img_folder)
    if module_found:
        process_bar = tqdm(total=all_files)
        process_bar.set_description('Processing:')
    for file in files:
        process_bar.update(1)
        img = img_folder + file.replace("txt", "jpg")
        shutil.copy(img, new_img_folder)

# create_img_folder()

B站小教程

提供该自制数据集用于YOLOv5训练的演示
https://www.bilibili.com/video/BV15r4y1u7T3?spm_id_from=333.999.0.0&vd_source=addcaa01f75ba41a03ef7559c37453ab

Github

可参考该结构组织文件夹
Github链接: https://github.com/Neo-YH/EasyToolforYourCustomDataset