【YOLOv10改进实战】**【3】YOLOv10“启动”——数据集的搭建:保姆级教程

【YOLOv10改进实战】**【3】YOLOv10“启动”🚀🚀——数据集的搭建:保姆级教程🥤🥤



提示:本文正式进入实战环节,喜欢的小伙伴请多多关注支持。本文仅供学习使用,创作不易,未经作者允许,不得搬运或转载!!!


前言🏝️🏝️

目标检测是计算机视觉领域的一个重要任务,它旨在识别图像中的目标对象,并确定它们的位置。当前目标检测在各领域已有广泛研究,每年都有大量的paper被发表。几乎所有的研究性论文都会使用数据集训练自己改进的模型,数据集是我们训练模型必不可缺的东西,也是paper发表的关键一步。
下文将介绍 数据集的获取及处理过程。

在这里插入图片描述


1. 数据集的获取🏝️🏝️

1.1 自建数据集
在我们的研究中,可能需要某些特定场景或高质量的数据,而现有公共数据集没有该类型数据或数据量很少,需要我们自己新建或扩充。可以考虑以下方法建立数据集:

  • 使用设备进行拍摄(如手机、相机等)
  • 可以从网站爬取自己需要的图篇数据,但需注意相关图片是否存在版权法律问题、隐私道德问题等。
  • 借助科技,使用合成算法(GAN)或AI工具生成图像。
    在这里插入图片描述

1.2公共数据集
对于公共数据集的获取,一般CSDN上好多up主的blog文章里或目标检测项目都有链接。虽然时大数据时代,每个人的信息都那啥,但大家点击链接时还是要仔细甄别。以下是一些提供公共数据集的网站,仅供参考:

  • GitHub: 这是一个最大最全面的开源网站,上面有大量的开源项目和数据集。
  • 阿里云天池数据集: 阿里云天池数据集平台汇集了官方、打榜、聚合、推荐、公共数据集。
  • Papers With Code: 这个平台不仅提供数据集,还有相关的模型、代码和论文。
  • OpenDataLab: 这是一个有影响力的数据开源开放平台,公开数据集触手可及。
  • Kaggle: Kaggle不仅是一个竞赛平台,也有海量的数据集资源。
  • OSF: 一个免费的开放平台,可以通过关键词搜索并获取数据集。
  • spm网站: 该网站提供了多模态面部数据集。
    在这里插入图片描述

2. 数据集标注🏝️🏝️

目标检测任务中,标注数据是一个重要的环节,它涉及到在图像中识别和定位感兴趣的目标。训练一个有效的目标检测模型,往往需要大量的标注数据,自建的数据集需要自己标注(“钞能力”除外)。标注数据集是包含目标对象的图像以及相应的标签信息的数据集,这些标签信息通常包括目标的类别和位置(通常是边界框)。

2.1 标注数据集对于目标检测模型训练至关重要的原因如下:

  • 监督学习:目标检测模型通常采用监督学习的方法进行训练。这意味着模型需要从带有正确标签的数据中学习如何识别和定位目标。标注数据集提供了这些必要的标签。
  • 特征学习:通过标注数据集,模型可以学习到不同目标对象的特征,包括形状、颜色、纹理等,这对于正确识别目标至关重要。
  • 定位精度:标注数据集中的边界框信息帮助模型学习如何精确地定位目标对象,这对于目标检测的准确性至关重要。
  • 泛化能力:使用多样化的标注数据集可以帮助模型学习到更广泛的特征,从而提高模型在不同场景和条件下的泛化能力。
  • 性能评估:标注数据集不仅用于训练,还用于评估目标检测模型的性能。通过比较模型预测的标签和实际的标注标签,可以评估模型的准确性、召回率等指标。
  • 数据增强:标注数据集可以用于数据增强,通过图像变换(如旋转、缩放、裁剪等)生成更多的训练样本,这有助于提高模型的鲁棒性。
  • 模型改进:标注数据集还可以用于模型的迭代改进。通过分析模型在标注数据集上的表现,可以识别模型的不足之处,并据此进行优化。
  • 多任务学习:在一些复杂的目标检测任务中,可能需要同时识别多个目标或执行其他视觉任务(如分割、关键点检测等)。标注数据集提供了执行这些多任务学习所需的信息。

2.2 常用的标注工具:

  • LabelImg:这是一个流行的开源图像标注工具,使用Python编写,图形界面基于Qt。它支持矩形框标注,并可以生成VOC和YOLO格式的标签文件。
  • Labelme:由MIT的CSAIL实验室开发,支持矩形、圆形、线段、点等多种形式的标注。它能够输出VOC和COCO格式的标注数据,并且支持视频标注。
  • RectLabel :支持对象检测和图像实例分割数据标注,可以导出YOLO、KITTI、COCOJSON和CSV格式。
  • OpenCV/CVAT :OpenCV提供的一个高效的计算机视觉标注工具,支持图像分类、对象检测、图像语义分割、实例分割数据标注。
  • VOTT (Microsoft’s Visual Object Tagging Tool):微软发布的基于WEB方式本地部署的视觉数据标注工具,支持多种格式的导出。
  • point-cloud-annotation-tool :针对3D点云数据标注的工具,支持点云数据加载、保存与可视化。
  • Boobs :一个专属的YOLO BBox标注工具,支持VOC/COCO格式数据导出,基于WEB方式的标注工具。
  • ImageLabelToolPlus :一款在线的深度学习图像分割标注工具,可以生成.json文件。
  • LabelWeb :一个基于网页的目标检测数据标注工具,方便团队合作标注。
  • KDAT (Kaggle Dataset Annotator Tool) :一个专为视觉方向目标检测、分割任务设计的数据标注工具。

2.3LabelImg安装和使用

本up主使用的是LabelImg,环境为Windows+Anaconda(3)。基于此,LabelImg的安装使用方法如下:
安装教程
打开 开始菜单栏,点击Anaconda3 -> Anaconda Prompt 进入命令行窗口:
在这里插入图片描述
在这里插入图片描述

创建一个新环境来安装labelimg,输入以下代码创建一个labelimg环境:


conda create -n labelimg  python=3.10  #labelimg为我们创建的文件名,可自己按喜好更改;
                                       #python=3.10为此环境python的版本,可自主选择,建议别如选太低

在这里插入图片描述
根据系统提示输入y或n(y代表确认,n代表取消)
在这里插入图片描述
输入以下代码检查环境是否安装成功:

conda info --envs #系统会展示安装的环境名

随后,输入以下代码激活环境:

conda activate labelimg #activate labelimg亦可

然后,在labelimg环境下继续输入以下代码安装labelimg工具,下面两行代码二选一:

conda install labelimg  #推荐用此方法安装
pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple  #不推荐,可能会出现依赖缺失或者版本不匹配问题 
#https://pypi.tuna.tsinghua.edu.cn/simple为清华镜像源,也可使用其他镜像源

安装完毕后直接输入以下代码打开labelimg

labelimg

labelimg界面如下所示:
在这里插入图片描述
如出现闪退等问题,需在labelimg环境下输入以下代码卸载labelimg,按上述步骤重新安装。

conda uninstall labelimg

使用教程
labelimg页面介绍如下(图源:不要洋洋洋):
在这里插入图片描述
YOLO格式数据(图源:不要洋洋洋):
在这里插入图片描述


3. 公共数据集处理 🏝️🏝️

由于YOLOv10使用的是YOLO格式的标注文件,公共数据了可能会涉及到格式转换问题,比如VOC格式转YOLO格式、YOLO格式转VOC格式 、CreateML格式转VOC格式。转换源码如下,可根据自己的需求选择。

CreateML (json) 2 voc(xml) 源码:

import os 
import json
from tqdm import tqdm
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--json_path', default='dataset/json',type=str, help="input: coco format(json)")#json文件路径
parser.add_argument('--save_path', default='dataset/txt', type=str, help="specify where to save the output dir of labels")#txt文件保存路径
arg = parser.parse_args()

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

    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

if __name__ == '__main__':
    json_file =   arg.json_path # COCO Object Instance 类型的标注
    ana_txt_save_path = arg.save_path  # 保存的路径

    data = json.load(open(json_file, 'r'))
    if not os.path.exists(ana_txt_save_path):
        os.makedirs(ana_txt_save_path)
    
    id_map = {} # coco数据集的id不连续!重新映射一下再输出!
    for i, category in enumerate(data['categories']): 
        id_map[category['id']] = i

    # 通过事先建表来降低时间复杂度
    max_id = 0
    for img in data['images']:
        max_id = max(max_id, img['id'])
    # 注意这里不能写作 [[]]*(max_id+1),否则列表内的空列表共享地址
    img_ann_dict = [[] for i in range(max_id+1)] 
    for i, ann in enumerate(data['annotations']):
        img_ann_dict[ann['image_id']].append(i)

    for img in tqdm(data['images']):
        filename = img["file_name"]
        img_width = img["width"]
        img_height = img["height"]
        img_id = img["id"]
        head, tail = os.path.splitext(filename)
        ana_txt_name = head + ".txt"  # 对应的txt名字,与jpg一致
        f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), 'w')
        '''for ann in data['annotations']:
            if ann['image_id'] == img_id:
                box = convert((img_width, img_height), ann["bbox"])
                f_txt.write("%s %s %s %s %s\n" % (id_map[ann["category_id"]], box[0], box[1], box[2], box[3]))'''
        # 这里可以直接查表而无需重复遍历
        for ann_id in img_ann_dict[img_id]:
            ann = data['annotations'][ann_id]
            box = convert((img_width, img_height), ann["bbox"])
            f_txt.write("%s %s %s %s %s\n" % (id_map[ann["category_id"]], box[0], box[1], box[2], box[3]))
        f_txt.close()

voc(xml) 2 yolo (txt)源码:

#xml 2 yolo
import xml.etree.ElementTree as ET
import os
from os import getcwd
import glob

# 1.
# 自己创建文件夹,例如:label_mal label_txt  也可以修改别的
image_set = 'datasets/tea/labels/val_xml_enhance_2'  # 需要转换的文件夹名称(文件夹内放xml标签文件)
imageset2 = 'datasets/tea/labels/val_txt_enhance_2'  # 保存txt的文件夹

# 2.
# 换成你的类别 当前的顺序,就txt 0,1,2,3 四个类别
classes = ['tea']  # 标注时的标签 注意顺序一定不要错。

# 3.
# # 转换文件夹的绝对路径
# data_dir = 'D:/detectAuto_/data'
# 或者 读取当前路径
data_dir = getcwd()  # 当前路径
'''
xml中框的左上角坐标和右下角坐标(x1,y1,x2,y2)
》》txt中的中心点坐标和宽和高(x,y,w,h),并且归一化
'''
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    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(data_dir, imageset1, imageset2, image_id):
    in_file = open(data_dir + '/%s/%s.xml' % (imageset1, image_id), encoding='UTF-8')  # 读取xml
    out_file = open(data_dir + '/%s/%s.txt' % (imageset2, image_id), 'w', encoding='UTF-8')  # 保存txt

    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        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)
        out_file.write(str(cls_id) + " " + " ".join([str('%.6f' % a) for a in bb]) + '\n')
image_ids = []
for x in glob.glob(data_dir + '/%s' % image_set + '/*.xml'):
    image_ids.append(os.path.basename(x)[:-4])
print('\n%s数量:' % image_set, len(image_ids))  # 确认数量
i = 0
for image_id in image_ids:
    i = i + 1
    convert_annotation(data_dir, image_set, imageset2, image_id)
    print("%s 数据:%s/%s文件完成!" % (image_set, i, len(image_ids)))

print("Done!!!")

yolo (txt) 2 voc(xml) 源码:

import os
import cv2
from tqdm import tqdm
from lxml.etree import Element, SubElement, tostring, ElementTree
from xml.dom.minidom import parseString
import numpy as np

out_root = r'datasets/tea/labels' # xml文件存放路径
def build_dir(out_dir):
    if not os.path.exists(out_dir):
        os.mkdir(out_dir)
    return out_dir
def get_root_lst(root, suffix='jpg', suffix_n=3):
    root_lst, name_lst = [], []

    for dir, file, names in os.walk(root):
        root_lst = root_lst + [os.path.join(dir, name) for name in names if name[-suffix_n:] == suffix]
        name_lst = name_lst + [name for name in names if name[-suffix_n:] == suffix]

    return root_lst, name_lst

def read_txt(path):
    txt_info_lst = []
    with open(path, "r", encoding='utf-8') as f:
        for line in f:
            txt_info_lst.append(list(line.strip('\n').split()))
    txt_info_lst = np.array(txt_info_lst)
    return txt_info_lst

def product_xml(name_img, boxes, codes, img=None, wh=None):
    '''
    :param img: 以读好的图片
    :param name_img: 图片名字
    :param boxes: box为列表
    :param codes: 为列表
    :return:
    '''
    if img is not None:
        width = img.shape[0]
        height = img.shape[1]
    else:
        assert wh is not None
        width = wh[0]
        height = wh[1]
    # print('xml w:{} h:{}'.format(width,height))

    node_root = Element('annotation')
    node_folder = SubElement(node_root, 'folder')
    node_folder.text = 'VOC2007'

    node_filename = SubElement(node_root, 'filename')
    node_filename.text = name_img  # 图片名字

    node_size = SubElement(node_root, 'size')
    node_width = SubElement(node_size, 'width')
    node_width.text = str(width)

    node_height = SubElement(node_size, 'height')
    node_height.text = str(height)

    node_depth = SubElement(node_size, 'depth')
    node_depth.text = '3'

    for i, code in enumerate(codes):
        box = [boxes[i][0], boxes[i][1], boxes[i][2], boxes[i][3]]
        node_object = SubElement(node_root, 'object')
        node_name = SubElement(node_object, 'name')
        node_name.text = code
        node_difficult = SubElement(node_object, 'difficult')
        node_difficult.text = '0'
        node_bndbox = SubElement(node_object, 'bndbox')
        node_xmin = SubElement(node_bndbox, 'xmin')
        node_xmin.text = str(int(box[0]))
        node_ymin = SubElement(node_bndbox, 'ymin')
        node_ymin.text = str(int(box[1]))
        node_xmax = SubElement(node_bndbox, 'xmax')
        node_xmax.text = str(int(box[2]))
        node_ymax = SubElement(node_bndbox, 'ymax')
        node_ymax.text = str(int(box[3]))

    xml = tostring(node_root, pretty_print=True)  # 格式化显示,该换行的换行
    dom = parseString(xml)

    name = name_img[:-4] + '.xml'

    tree = ElementTree(node_root)

    # print('name:{},dom:{}'.format(name, dom))
    return tree, name


def yolov5txt2xml(root_data, txt_root, gt_labels=None, out_dir=None):
    # 获得图像与txt的路径与名称的列表
    img_roots_lst, img_names_lst = get_root_lst(root_data, suffix='jpg', suffix_n=3)
    txt_roots_lst, txt_names_lst = get_root_lst(txt_root, suffix='txt', suffix_n=3)
    # 创建保存xml的文件
    out_dir = build_dir(out_dir) if out_dir is not None else build_dir(os.path.join(out_root, 'val_xml')) # val_xml xml文件存放的文件夹

    label_str_lst = []
    # 通过图像遍历
    for i, img_root in tqdm(enumerate(img_roots_lst)):
        # 获得图像名称,并得到对应txt名称
        img_name = img_names_lst[i]
        txt_name = img_name[:-3] + 'txt'

        if txt_name in txt_names_lst:  # 通过图像获得txt名称是否存在,存在则继续,否则不继续
            txt_index = list(txt_names_lst).index(str(txt_name))  # 获得列表txt对应索引,以便后续获得路径
            # 通过图像获得图像高与宽
            img = cv2.imread(img_root)
            height, width = img.shape[:2]
            # 读取对应txt的信息
            txt_info = read_txt(txt_roots_lst[txt_index])

            # 以下获得txt信息,并保存labels_lst与boxes_lst中,且一一对应
            labels_lst, boxes_lst = [], []
            for info in txt_info:
                label_str = str(info[0])
                if label_str not in label_str_lst:
                    label_str_lst.append(label_str)

                x, y, w, h = float(info[1]) * width, float(info[2]) * height, float(info[3]) * width, float(
                    info[4]) * height
                xmin, ymin, xmax, ymax = int(x - w / 2), int(y - h / 2), int(x + w / 2), int(y + h / 2)
                labels_lst.append(label_str)
                boxes_lst.append([xmin, ymin, xmax, ymax])
            # 是否转换信息
            if gt_labels:  # gt_labels需要和txt类别对应
                labels_lst = [gt_labels[int(lb)] for lb in labels_lst]
            # 构建xml文件
            if len(labels_lst) > 0:
                tree, xml_name = product_xml(img_name, boxes_lst, labels_lst, wh=[w, h])
                tree.write(os.path.join(out_dir, xml_name))

    print('gt label:', gt_labels)
    print('txt label:', label_str_lst)
    print('save root:', out_dir)


if __name__ == '__main__':
    # Volumes/ACASIS_Media/DeepLearning/ultralytics/datasets/preson
    root_path = r'datasets/tea/images/val'  # 图片路径
    txt_root = r'datasets/tea/labels/val'   # txt路径

    gt_labels = ['tea']

    yolov5txt2xml(root_path, txt_root, gt_labels=gt_labels)


4. 数据增强🏝️🏝️

由于深度学习需要大量数据集,如果数据量很少,可能会使模型出现过拟合、泛化能力差、性能受限、评估困难等问题。可以使用数据增强、主动学习、迁移学习,交叉训练等方式缓解这一问题。 数据增强(Data Augmentation): 通过对现有数据进行变换(如旋转、翻转、缩放、裁剪等)来增加数据集的多样性。
示例:

在这里插入图片描述
数据增强源码,使用的是原图+voc格式文件:

import time
import random
import copy
import cv2
import os
import math
import numpy as np
from skimage.util import random_noise
from lxml import etree, objectify
import xml.etree.ElementTree as ET
import argparse


# 显示图片
def show_pic(img, bboxes=None):
    '''
    输入:
        img:图像array
        bboxes:图像的所有boudning box list, 格式为[[x_min, y_min, x_max, y_max]....]
        names:每个box对应的名称
    '''
    for i in range(len(bboxes)):
        bbox = bboxes[i]
        x_min = bbox[0]
        y_min = bbox[1]
        x_max = bbox[2]
        y_max = bbox[3]
        cv2.rectangle(img, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 3)
    cv2.namedWindow('pic', 0)  # 1表示原图
    cv2.moveWindow('pic', 0, 0)
    cv2.resizeWindow('pic', 1200, 800)  # 可视化的图片大小
    cv2.imshow('pic', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    
# 图像均为cv2读取
class DataAugmentForObjectDetection():
    def __init__(self, rotation_rate=0.5, max_rotation_angle=5,
                 crop_rate=0.5, shift_rate=0.5, change_light_rate=0.5,
                 add_noise_rate=0.5, flip_rate=0.5,
                 cutout_rate=0.5, cut_out_length=50, cut_out_holes=1, cut_out_threshold=0.5,
                 is_addNoise=True, is_changeLight=True, is_cutout=True, is_rotate_img_bbox=True,
                 is_crop_img_bboxes=True, is_shift_pic_bboxes=True, is_filp_pic_bboxes=True):

        # 配置各个操作的属性
        self.rotation_rate = rotation_rate
        self.max_rotation_angle = max_rotation_angle
        self.crop_rate = crop_rate
        self.shift_rate = shift_rate
        self.change_light_rate = change_light_rate
        self.add_noise_rate = add_noise_rate
        self.flip_rate = flip_rate
        self.cutout_rate = cutout_rate

        self.cut_out_length = cut_out_length
        self.cut_out_holes = cut_out_holes
        self.cut_out_threshold = cut_out_threshold

        # 是否使用某种增强方式
        self.is_addNoise = is_addNoise
        self.is_changeLight = is_changeLight
        self.is_cutout = is_cutout
        self.is_rotate_img_bbox = is_rotate_img_bbox
        self.is_crop_img_bboxes = is_crop_img_bboxes
        self.is_shift_pic_bboxes = is_shift_pic_bboxes
        self.is_filp_pic_bboxes = is_filp_pic_bboxes

    # ----1.加噪声---- #
    def _addNoise(self, img):
        '''
        输入:
            img:图像array
        输出:
            加噪声后的图像array,由于输出的像素是在[0,1]之间,所以得乘以255
        '''
        # return cv2.GaussianBlur(img, (11, 11), 0)
        return random_noise(img, mode='gaussian', seed=int(time.time()), clip=True) * 255

    # ---2.调整亮度--- #
    def _changeLight(self, img):
        alpha = random.uniform(0.35, 1)
        blank = np.zeros(img.shape, img.dtype)
        return cv2.addWeighted(img, alpha, blank, 1 - alpha, 0)

    # ---3.cutout--- #
    def _cutout(self, img, bboxes, length=100, n_holes=1, threshold=0.5):
        '''
        原版本:https://github.com/uoguelph-mlrg/Cutout/blob/master/util/cutout.py
        Randomly mask out one or more patches from an image.
        Args:
            img : a 3D numpy array,(h,w,c)
            bboxes : 框的坐标
            n_holes (int): Number of patches to cut out of each image.
            length (int): The length (in pixels) of each square patch.
        '''

        def cal_iou(boxA, boxB):
            '''
            boxA, boxB为两个框,返回iou
            boxB为bouding box
            '''
            # determine the (x, y)-coordinates of the intersection rectangle
            xA = max(boxA[0], boxB[0])
            yA = max(boxA[1], boxB[1])
            xB = min(boxA[2], boxB[2])
            yB = min(boxA[3], boxB[3])

            if xB <= xA or yB <= yA:
                return 0.0

            # compute the area of intersection rectangle
            interArea = (xB - xA + 1) * (yB - yA + 1)

            # compute the area of both the prediction and ground-truth
            # rectangles
            boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
            boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
            iou = interArea / float(boxBArea)
            return iou

        # 得到h和w
        if img.ndim == 3:
            h, w, c = img.shape
        else:
            _, h, w, c = img.shape
        mask = np.ones((h, w, c), np.float32)
        for n in range(n_holes):
            chongdie = True  # 看切割的区域是否与box重叠太多
            while chongdie:
                y = np.random.randint(h)
                x = np.random.randint(w)

                y1 = np.clip(y - length // 2, 0,
                             h)  # numpy.clip(a, a_min, a_max, out=None), clip这个函数将将数组中的元素限制在a_min, a_max之间,大于a_max的就使得它等于 a_max,小于a_min,的就使得它等于a_min
                y2 = np.clip(y + length // 2, 0, h)
                x1 = np.clip(x - length // 2, 0, w)
                x2 = np.clip(x + length // 2, 0, w)

                chongdie = False
                for box in bboxes:
                    if cal_iou([x1, y1, x2, y2], box) > threshold:
                        chongdie = True
                        break
            mask[y1: y2, x1: x2, :] = 0.
        img = img * mask
        return img

    # ---4.旋转--- #
    def _rotate_img_bbox(self, img, bboxes, angle=5, scale=1.):
        '''
        参考:https://blog.csdn.net/u014540717/article/details/53301195crop_rate
        输入:
            img:图像array,(h,w,c)
            bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
            angle:旋转角度
            scale:默认1
        输出:
            rot_img:旋转后的图像array
            rot_bboxes:旋转后的boundingbox坐标list
        '''
        # 旋转图像
        w = img.shape[1]
        h = img.shape[0]
        # 角度变弧度
        rangle = np.deg2rad(angle)  # angle in radians
        # now calculate new image width and height
        nw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)) * scale
        nh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)) * scale
        # ask OpenCV for the rotation matrix
        rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)
        # calculate the move from the old center to the new center combined
        # with the rotation
        rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))
        # the move only affects the translation, so update the translation
        rot_mat[0, 2] += rot_move[0]
        rot_mat[1, 2] += rot_move[1]
        # 仿射变换
        rot_img = cv2.warpAffine(img, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))), flags=cv2.INTER_LANCZOS4)

        # 矫正bbox坐标
        # rot_mat是最终的旋转矩阵
        # 获取原始bbox的四个中点,然后将这四个点转换到旋转后的坐标系下
        rot_bboxes = list()
        for bbox in bboxes:
            xmin = bbox[0]
            ymin = bbox[1]
            xmax = bbox[2]
            ymax = bbox[3]
            point1 = np.dot(rot_mat, np.array([(xmin + xmax) / 2, ymin, 1]))
            point2 = np.dot(rot_mat, np.array([xmax, (ymin + ymax) / 2, 1]))
            point3 = np.dot(rot_mat, np.array([(xmin + xmax) / 2, ymax, 1]))
            point4 = np.dot(rot_mat, np.array([xmin, (ymin + ymax) / 2, 1]))
            # 合并np.array
            concat = np.vstack((point1, point2, point3, point4))
            # 改变array类型
            concat = concat.astype(np.int32)
            # 得到旋转后的坐标
            rx, ry, rw, rh = cv2.boundingRect(concat)
            rx_min = rx
            ry_min = ry
            rx_max = rx + rw
            ry_max = ry + rh
            # 加入list中
            rot_bboxes.append([rx_min, ry_min, rx_max, ry_max])

        return rot_img, rot_bboxes

    # ---5.裁剪--- #
    def _crop_img_bboxes(self, img, bboxes):
        '''
        裁剪后的图片要包含所有的框
        输入:
            img:图像array
            bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
        输出:
            crop_img:裁剪后的图像array
            crop_bboxes:裁剪后的bounding box的坐标list
        '''
        # 裁剪图像
        w = img.shape[1]
        h = img.shape[0]
        x_min = w  # 裁剪后的包含所有目标框的最小的框
        x_max = 0
        y_min = h
        y_max = 0
        for bbox in bboxes:
            x_min = min(x_min, bbox[0])
            y_min = min(y_min, bbox[1])
            x_max = max(x_max, bbox[2])
            y_max = max(y_max, bbox[3])

        d_to_left = x_min  # 包含所有目标框的最小框到左边的距离
        d_to_right = w - x_max  # 包含所有目标框的最小框到右边的距离
        d_to_top = y_min  # 包含所有目标框的最小框到顶端的距离
        d_to_bottom = h - y_max  # 包含所有目标框的最小框到底部的距离

        # 随机扩展这个最小框
        crop_x_min = int(x_min - random.uniform(0, d_to_left))
        crop_y_min = int(y_min - random.uniform(0, d_to_top))
        crop_x_max = int(x_max + random.uniform(0, d_to_right))
        crop_y_max = int(y_max + random.uniform(0, d_to_bottom))

        # 随机扩展这个最小框 , 防止别裁的太小
        # crop_x_min = int(x_min - random.uniform(d_to_left//2, d_to_left))
        # crop_y_min = int(y_min - random.uniform(d_to_top//2, d_to_top))
        # crop_x_max = int(x_max + random.uniform(d_to_right//2, d_to_right))
        # crop_y_max = int(y_max + random.uniform(d_to_bottom//2, d_to_bottom))

        # 确保不要越界
        crop_x_min = max(0, crop_x_min)
        crop_y_min = max(0, crop_y_min)
        crop_x_max = min(w, crop_x_max)
        crop_y_max = min(h, crop_y_max)

        crop_img = img[crop_y_min:crop_y_max, crop_x_min:crop_x_max]

        # 裁剪boundingbox
        # 裁剪后的boundingbox坐标计算
        crop_bboxes = list()
        for bbox in bboxes:
            crop_bboxes.append([bbox[0] - crop_x_min, bbox[1] - crop_y_min, bbox[2] - crop_x_min, bbox[3] - crop_y_min])

        return crop_img, crop_bboxes

    # ---6.平移--- #
    def _shift_pic_bboxes(self, img, bboxes):
        '''
        平移后的图片要包含所有的框
        输入:
            img:图像array
            bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
        输出:
            shift_img:平移后的图像array
            shift_bboxes:平移后的bounding box的坐标list
        '''
        # 平移图像
        w = img.shape[1]
        h = img.shape[0]
        x_min = w  # 裁剪后的包含所有目标框的最小的框
        x_max = 0
        y_min = h
        y_max = 0
        for bbox in bboxes:
            x_min = min(x_min, bbox[0])
            y_min = min(y_min, bbox[1])
            x_max = max(x_max, bbox[2])
            y_max = max(y_max, bbox[3])

        d_to_left = x_min  # 包含所有目标框的最大左移动距离
        d_to_right = w - x_max  # 包含所有目标框的最大右移动距离
        d_to_top = y_min  # 包含所有目标框的最大上移动距离
        d_to_bottom = h - y_max  # 包含所有目标框的最大下移动距离

        x = random.uniform(-(d_to_left - 1) / 3, (d_to_right - 1) / 3)
        y = random.uniform(-(d_to_top - 1) / 3, (d_to_bottom - 1) / 3)

        M = np.float32([[1, 0, x], [0, 1, y]])  # x为向左或右移动的像素值,正为向右负为向左; y为向上或者向下移动的像素值,正为向下负为向上
        shift_img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))

        #  平移boundingbox
        shift_bboxes = list()
        for bbox in bboxes:
            shift_bboxes.append([bbox[0] + x, bbox[1] + y, bbox[2] + x, bbox[3] + y])

        return shift_img, shift_bboxes

    # ---7.镜像--- #
    def _filp_pic_bboxes(self, img, bboxes):
        '''
            平移后的图片要包含所有的框
            输入:
                img:图像array
                bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
            输出:
                flip_img:平移后的图像array
                flip_bboxes:平移后的bounding box的坐标list
        '''
        # 翻转图像

        flip_img = copy.deepcopy(img)
        h, w, _ = img.shape

        sed = random.random()

        if 0 < sed < 0:  # 0.33的概率水平翻转,0.33的概率垂直翻转,0.33是对角反转
            flip_img = cv2.flip(flip_img, 0)  # _flip_x
            inver = 0
        elif 0 <= sed <= 1:
            flip_img = cv2.flip(flip_img, 1)  # _flip_y
            inver = 1
        else:
            flip_img = cv2.flip(flip_img, -1)  # flip_x_y
            inver = -1

        # 调整boundingbox
        flip_bboxes = list()
        for box in bboxes:
            x_min = box[0]
            y_min = box[1]
            x_max = box[2]
            y_max = box[3]

            if inver == 0:
                # 0:垂直翻转
                flip_bboxes.append([x_min, h - y_max, x_max, h - y_min])
            elif inver == 1:
                # 1:水平翻转
                flip_bboxes.append([w - x_max, y_min, w - x_min, y_max])
            elif inver == -1:
                # -1:水平垂直翻转
                flip_bboxes.append([w - x_max, h - y_max, w - x_min, h - y_min])
        return flip_img, flip_bboxes

    # 图像增强方法
    def dataAugment(self, img, bboxes):
        '''
        图像增强
        输入:
            img:图像array
            bboxes:该图像的所有框坐标
        输出:
            img:增强后的图像
            bboxes:增强后图片对应的box
        '''
        change_num = 0  # 改变的次数
        # print('------')
        while change_num < 1:  # 默认至少有一种数据增强生效

            if self.is_rotate_img_bbox:
                if random.random() > self.rotation_rate:  # 旋转
                    change_num += 1
                    angle = random.uniform(-self.max_rotation_angle, self.max_rotation_angle)
                    scale = random.uniform(0.7, 0.8)
                    img, bboxes = self._rotate_img_bbox(img, bboxes, angle, scale)

            if self.is_shift_pic_bboxes:
                if random.random() < self.shift_rate:  # 平移
                    change_num += 1
                    img, bboxes = self._shift_pic_bboxes(img, bboxes)

            if self.is_changeLight:
                if random.random() > self.change_light_rate:  # 改变亮度
                    change_num += 1
                    img = self._changeLight(img)

            if self.is_addNoise:
                if random.random() < self.add_noise_rate:  # 加噪声
                    change_num += 1
                    img = self._addNoise(img)
            if self.is_cutout:
                if random.random() < self.cutout_rate:  # cutout
                    change_num += 1
                    img = self._cutout(img, bboxes, length=self.cut_out_length, n_holes=self.cut_out_holes,
                                       threshold=self.cut_out_threshold)
            if self.is_filp_pic_bboxes:
                if random.random() < self.flip_rate:  # 翻转
                    change_num += 1
                    img, bboxes = self._filp_pic_bboxes(img, bboxes)

        return img, bboxes


# xml解析工具
class ToolHelper():
    # 从xml文件中提取bounding box信息, 格式为[[x_min, y_min, x_max, y_max, name]]
    def parse_xml(self, path):
        '''
        输入:
            xml_path: xml的文件路径
        输出:
            从xml文件中提取bounding box信息, 格式为[[x_min, y_min, x_max, y_max, name]]
        '''
        tree = ET.parse(path)
        root = tree.getroot()
        objs = root.findall('object')
        coords = list()
        for ix, obj in enumerate(objs):
            name = obj.find('name').text
            box = obj.find('bndbox')
            x_min = int(box[0].text)
            y_min = int(box[1].text)
            x_max = int(box[2].text)
            y_max = int(box[3].text)
            coords.append([x_min, y_min, x_max, y_max, name])
        return coords

    # 保存图片结果
    def save_img(self, file_name, save_folder, img):
        cv2.imwrite(os.path.join(save_folder, file_name), img)

    # 保持xml结果
    def save_xml(self, file_name, save_folder, img_info, height, width, channel, bboxs_info):
        '''
        :param file_name:文件名
        :param save_folder:#保存的xml文件的结果
        :param height:图片的信息
        :param width:图片的宽度
        :param channel:通道
        :return:
        '''
        folder_name, img_name = img_info  # 得到图片的信息

        E = objectify.ElementMaker(annotate=False)

        anno_tree = E.annotation(
            E.folder(folder_name),
            E.filename(img_name),
            E.path(os.path.join(folder_name, img_name)),
            E.source(
                E.database('Unknown'),
            ),
            E.size(
                E.width(width),
                E.height(height),
                E.depth(channel)
            ),
            E.segmented(0),
        )

        labels, bboxs = bboxs_info  # 得到边框和标签信息
        for label, box in zip(labels, bboxs):
            anno_tree.append(
                E.object(
                    E.name(label),
                    E.pose('Unspecified'),
                    E.truncated('0'),
                    E.difficult('0'),
                    E.bndbox(
                        E.xmin(box[0]),
                        E.ymin(box[1]),
                        E.xmax(box[2]),
                        E.ymax(box[3])
                    )
                ))

        etree.ElementTree(anno_tree).write(os.path.join(save_folder, file_name), pretty_print=True)


if __name__ == '__main__':

    need_aug_num = 5  # 每张图片需要增强的次数

    is_endwidth_dot = True  # 文件是否以.jpg或者png结尾

    dataAug = DataAugmentForObjectDetection()  # 数据增强工具类

    toolhelper = ToolHelper()  # 工具

    # 获取相关参数
    parser = argparse.ArgumentParser()
    parser.add_argument('--source_img_path', type=str, default='datasets/tea/images/val')# 图片原始位置
    parser.add_argument('--source_xml_path', type=str, default='datasets/tea/labels/val_xml')# xml的原始位置
    parser.add_argument('--save_img_path', type=str, default='datasets/tea/images/val_enhance_2')# 图片增强结果保存文件
    parser.add_argument('--save_xml_path', type=str, default='datasets/tea/labels/val_xml_enhance_2')# xml增强结果保存文件
    args = parser.parse_args()
    source_img_path = args.source_img_path  
    source_xml_path = args.source_xml_path  

    save_img_path = args.save_img_path  
    save_xml_path = args.save_xml_path  

    # 如果保存文件夹不存在就创建
    if not os.path.exists(save_img_path):
        os.mkdir(save_img_path)

    if not os.path.exists(save_xml_path):
        os.mkdir(save_xml_path)

    for parent, _, files in os.walk(source_img_path):
        files.sort()
        for file in files:
            cnt = 0
            pic_path = os.path.join(parent, file)
            xml_path = os.path.join(source_xml_path, file[:-4] + '.xml')
            values = toolhelper.parse_xml(xml_path)  # 解析得到box信息,格式为[[x_min,y_min,x_max,y_max,name]]
            coords = [v[:4] for v in values]  # 得到框
            labels = [v[-1] for v in values]  # 对象的标签

            # 如果图片是有后缀的
            if is_endwidth_dot:
                # 找到文件的最后名字
                dot_index = file.rfind('.')
                _file_prefix = file[:dot_index]  # 文件名的前缀
                _file_suffix = file[dot_index:]  # 文件名的后缀
            img = cv2.imread(pic_path)

            # show_pic(img, coords)  # 显示原图
            while cnt < need_aug_num:  # 继续增强
                auged_img, auged_bboxes = dataAug.dataAugment(img, coords)
                auged_bboxes_int = np.array(auged_bboxes).astype(np.int32)
                height, width, channel = auged_img.shape  # 得到图片的属性
                img_name = '{}_{}{}'.format(_file_prefix, cnt + 1, _file_suffix)  # 图片保存的信息
                toolhelper.save_img(img_name, save_img_path,
                                    auged_img)  # 保存增强图片

                toolhelper.save_xml('{}_{}.xml'.format(_file_prefix, cnt + 1),
                                    save_xml_path, (save_img_path, img_name), height, width, channel,
                                    (labels, auged_bboxes_int))  # 保存xml文件
                # show_pic(auged_img, auged_bboxes)  # 强化后的图
                print(img_name)
                cnt += 1  # 继续增强下一张


5. 划分数据集🏝️🏝️

在目标检测任务中,合理划分训练集、验证集和测试集对于评估模型性能和避免过拟合至关重要。数据集划分没有固定规则,需要注意的是,划分比例应根据具体任务、数据集的大小和多样性以及模型的复杂性来确定。以下是常见的划分方法:

  • 60%训练集 / 20%验证集 /20%测试集:这是一种常见的划分比例,它为模型提供了足够的数据进行训练,同时保留了相当数量的数据用于模型验证和测试。
  • 70%训练集 / 15%验证集 / 15%测试集:这种划分给予训练集更多的数据,有助于模型更好地学习特征,同时保持一定量的验证和测试数据。
  • 80%训练集 / 10%验证集 / 10%测试集:当数据集较大时,这种划分可以提供大量的数据供模型训练,同时保证有适量的数据用于评估。
  • 90%训练集 / 5%验证集 /5%测试集:在数据集非常大的情况下,这种划分可以最大化训练数据量,但由于验证集和测试集较小,可能不足以准确评估模型的泛化能力。
  • 分层抽样(Stratified Sampling):在某些情况下,可能还需要考虑数据的分布。分层抽样确保每个类别在训练集、验证集和测试集中的比例与原始数据集中的比例相同,这有助于模型学习类别的分布。
  • 交叉验证(Cross-Validation):在数据量非常有限的情况下,可以使用交叉验证来更有效地利用数据。这种方法将数据集分成几个小的折叠,每个折叠轮流作为验证集,而其余部分作为训练集。
  • 动态数据集划分:在某些情况下,可以根据模型在验证集上的表现动态调整训练集和验证集的大小,以优化模型性能。

划分数据集源码:

# 将图片和标注数据按比例切分为 训练集和测试集
import shutil
import random
import os
 
# 原始路径
image_original_path = "dataset/JPEGImages/"
label_original_path = "dataset/txt/"
 
cur_path = os.getcwd()
 
# 训练集路径
train_image_path = os.path.join(cur_path, "dataset/images/train/")
train_label_path = os.path.join(cur_path, "dataset/labels/train/")
 
# 验证集路径
val_image_path = os.path.join(cur_path, "dataset/images/val/")
val_label_path = os.path.join(cur_path, "dataset/labels/val/")
 
# 测试集路径
test_image_path = os.path.join(cur_path, "dataset/images/test/")
test_label_path = os.path.join(cur_path, "dataset/labels/test/")
 
# 训练集、验证集、测试集目录
list_train = os.path.join(cur_path, "dataset//train.txt")
list_val = os.path.join(cur_path, "dataset//val.txt")
list_test = os.path.join(cur_path, "dataset//test.txt")
 
#划分比例
train_percent = 0.8
val_percent = 0.1
test_percent = 0.1
 
 
def del_file(path):
    for i in os.listdir(path):
        file_data = path + "\\" + i
        os.remove(file_data)
 
 
def mkdir():
    if not os.path.exists(train_image_path):
        os.makedirs(train_image_path)
    else:
        del_file(train_image_path)
    if not os.path.exists(train_label_path):
        os.makedirs(train_label_path)
    else:
        del_file(train_label_path)
 
    if not os.path.exists(val_image_path):
        os.makedirs(val_image_path)
    else:
        del_file(val_image_path)
    if not os.path.exists(val_label_path):
        os.makedirs(val_label_path)
    else:
        del_file(val_label_path)
 
    if not os.path.exists(test_image_path):
        os.makedirs(test_image_path)
    else:
        del_file(test_image_path)
    if not os.path.exists(test_label_path):
        os.makedirs(test_label_path)
    else:
        del_file(test_label_path)
 
 
def clearfile():
    if os.path.exists(list_train):
        os.remove(list_train)
    if os.path.exists(list_val):
        os.remove(list_val)
    if os.path.exists(list_test):
        os.remove(list_test)
 
 
def main():
    mkdir()
    clearfile()
 
    file_train = open(list_train, 'w')
    file_val = open(list_val, 'w')
    file_test = open(list_test, 'w')
 
    total_txt = os.listdir(label_original_path)
    num_txt = len(total_txt)
    list_all_txt = range(num_txt)
 
    num_train = int(num_txt * train_percent)
    num_val = int(num_txt * val_percent)
    num_test = num_txt - num_train - num_val
 
    train = random.sample(list_all_txt, num_train)
    # train从list_all_txt取出num_train个元素
    # 所以list_all_txt列表只剩下了这些元素
    val_test = [i for i in list_all_txt if not i in train]
    # 再从val_test取出num_val个元素,val_test剩下的元素就是test
    val = random.sample(val_test, num_val)
 
    print("训练集数目:{}, 验证集数目:{}, 测试集数目:{}".format(len(train), len(val), len(val_test) - len(val)))
    for i in list_all_txt:
        name = total_txt[i][:-4]
 
        srcImage = image_original_path + name + '.jpg'
        srcLabel = label_original_path + name + ".txt"
 
        if i in train:
            dst_train_Image = train_image_path + name + '.jpg'
            dst_train_Label = train_label_path + name + '.txt'
            shutil.copyfile(srcImage, dst_train_Image)
            shutil.copyfile(srcLabel, dst_train_Label)
            file_train.write(dst_train_Image + '\n')
        elif i in val:
            dst_val_Image = val_image_path + name + '.jpg'
            dst_val_Label = val_label_path + name + '.txt'
            shutil.copyfile(srcImage, dst_val_Image)
            shutil.copyfile(srcLabel, dst_val_Label)
            file_val.write(dst_val_Image + '\n')
        else:
            dst_test_Image = test_image_path + name + '.jpg'
            dst_test_Label = test_label_path + name + '.txt'
            shutil.copyfile(srcImage, dst_test_Image)
            shutil.copyfile(srcLabel, dst_test_Label)
            file_test.write(dst_test_Image + '\n')
 
    file_train.close()
    file_val.close()
    file_test.close()
 
 
if __name__ == "__main__":
    main()

参考资料🏝️🏝️

[1]https://gitcode.com/384520505/YoloDatasetsEnhance/blob/main/enhance_engine.py


本文内容到此结束,文章持续更新中,敬请期待!!
请添加图片描述
下一篇:YOLOv10训练自己的数据集:保姆级教程🥤🥤(预告)


  • 33
    点赞
  • 29
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值