【YOLOV5-5.x 源码解读】datasets.py
2024-04-02 13:40:31
目录
- 前言
- 0、导入需要的包和基本配置
- 1、相机设置
- 2、create_dataloader
- 3、自定义DataLoader
- 4、LoadImagesAndLabels
- 4.1、__init__
- 4.2、cache_labels
- 4.3、__len__
- 4.4.、__getitem__
- 4.5、collate_fn
- 4.6、collate_fn4
- 5、img2label_paths
- 6、verify_image_label
- 7、load_image
- 8、augment_hsv
- 9、load_mosaic、load_mosaic9
- 9.1、load_mosaic
- 9.2、load_mosaic9
- 10、random_perspective
- 11、box_candidates
- 12、replicate
- 13、letterbox
- 14、cutout
- 15、mixup
- 16、LoadImages、LoadStreams、LoadWebcam
- 17、hist_equalize
- 18、create_folder
- 19、flatten_recursive
- 20、extract_boxes
- 21、autosplit
- 22、dataset_stats
- 总结
前言
源码： YOLOv5源码.
导航: 【YOLOV5-5.x 源码讲解】整体项目文件导航.
注释版全部项目文件已上传至GitHub: yolov5-5.x-annotations.
这个文件主要是进行数据增强操作。
0、导入需要的包和基本配置
import glob # python自己带的一个文件操作相关模块 查找符合自己目的的文件(如模糊匹配)
import hashlib # 哈希模块 提供了多种安全方便的hash方法
import json # json文件操作模块
import logging # 日志模块
import math # 数学公式模块
import os # 与操作系统进行交互的模块 包含文件路径操作和解析
import random # 生成随机数模块
import shutil # 文件夹、压缩包处理模块
import time # 时间模块 更底层
from itertools import repeat # 复制模块
from multiprocessing.pool import ThreadPool, Pool # 多线程模块 线程池
from pathlib import Path # Path将str转换为Path对象 使字符串路径易于操作的模块
from threading import Thread # 多线程操作模块import cv2 # opencv模块
import numpy as np # numpy矩阵操作模块
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt # matplotlib画图模块
import torch # PyTorch深度学习模块
import torch.nn.functional as F # PyTorch函数接口 封装了很多卷积、池化等函数
import yaml # yaml文件操作模块
from PIL import Image, ExifTags # 图片、相机操作模块
from torch.utils.data import Dataset # 自定义数据集模块
from tqdm import tqdm # 进度条模块from utils.general import check_requirements, check_file, check_dataset, xywh2xyxy, xywhn2xyxy, xyxy2xywhn, \xyn2xy, segment2box, segments2boxes, resample_segments, clean_str
from utils.torch_utils import torch_distributed_zero_first# Parameters
help_url = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
img_formats = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo'] # acceptable image suffixes
vid_formats = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv'] # acceptable video suffixes
num_threads = min(8, os.cpu_count()) # 定义多线程个数
logger = logging.getLogger(__name__) # 初始化日志
1、相机设置
\qquad这部分是相机相关设置，当使用相机采样时才会使用。
# 相机设置
# Get orientation exif tag
# 专门为数码相机的照片而设定 可以记录数码照片的属性信息和拍摄数据
for orientation in ExifTags.TAGS.keys():if ExifTags.TAGS[orientation] == 'Orientation':break
def get_hash(paths):# 返回文件列表的hash值# Returns a single hash value of a list of paths (files or dirs)size = sum(os.path.getsize(p) for p in paths if os.path.exists(p)) # sizesh = hashlib.md5(str(size).encode()) # hash sizesh.update(''.join(paths).encode()) # hash pathsreturn h.hexdigest() # return hash
def exif_size(img):# 获取数码相机的图片宽高信息 并且判断是否需要旋转（数码相机可以多角度拍摄）# Returns exif-corrected PIL sizes = img.size # (width, height)try:rotation = dict(img._getexif().items())[orientation]if rotation == 6: # rotation 270s = (s[1], s[0])elif rotation == 8: # rotation 90s = (s[1], s[0])except:passreturn s
2、create_dataloader
\qquad自定义dataloader函数: 调用LoadImagesAndLabels获取数据集dataset(包括数据增强) + 调用分布式采样器DistributedSampler + 自定义InfiniteDataLoader 进行永久持续的采样数据 + 获取dataloader。关键核心是LoadImagesAndLabels()，这个文件的后面所有代码都是围绕这个模块进行的。
create_dataloader模块代码：
def create_dataloader(path, imgsz, batch_size, stride, single_cls=False,hyp=None, augment=False, cache=False, pad=0.0, rect=False,rank=-1, workers=8, image_weights=False, quad=False, prefix=''):"""在train.py中被调用，用于生成Trainloader, dataset，testloader自定义dataloader函数: 调用LoadImagesAndLabels获取数据集(包括数据增强) + 调用分布式采样器DistributedSampler +自定义InfiniteDataLoader 进行永久持续的采样数据:param path: 图片数据加载路径 train/test 如: ../datasets/VOC/images/train2007:param imgsz: train/test图片尺寸（数据增强后大小） 640:param batch_size: batch size 大小 8/16/32:param stride: 模型最大stride=32 [32 16 8]:param single_cls: 数据集是否是单类别 默认False:param hyp: 超参列表dict 网络训练时的一些超参数，包括学习率等，这里主要用到里面一些关于数据增强(旋转、平移等)的系数:param augment: 是否要进行数据增强 True:param cache: 是否cache_images False:param pad: 设置矩形训练的shape时进行的填充 默认0.0:param rect: 是否开启矩形train/test 默认训练集关闭 验证集开启:param rank: 多卡训练时的进程编号 rank为进程编号 -1且gpu=1时不进行分布式 -1且多块gpu使用DataParallel模式 默认-1:param workers: dataloader的numworks 加载数据时的cpu进程数:param image_weights: 训练时是否根据图片样本真实框分布权重来选择图片 默认False:param quad: dataloader取数据时, 是否使用collate_fn4代替collate_fn 默认False:param prefix: 显示信息 一个标志，多为train/val，处理标签时保存cache文件会用到"""# Make sure only the first process in DDP process the dataset first, and the following others can use the cache# 主进程实现数据的预读取并缓存，然后其它子进程则从缓存中读取数据并进行一系列运算。# 为了完成数据的正常同步, yolov5基于torch.distributed.barrier()函数实现了上下文管理器with torch_distributed_zero_first(rank):# 载入文件数据(增强数据集)dataset = LoadImagesAndLabels(path, imgsz, batch_size,augment=augment, # augment imageshyp=hyp, # augmentation hyperparametersrect=rect, # rectangular trainingcache_images=cache,single_cls=single_cls,stride=int(stride),pad=pad,image_weights=image_weights,prefix=prefix)batch_size = min(batch_size, len(dataset)) # bsnw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, workers]) # number of workers# 分布式采样器DistributedSamplersampler = torch.utils.data.distributed.DistributedSampler(dataset) if rank != -1 else None# 使用InfiniteDataLoader和_RepeatSampler来对DataLoader进行封装, 代替原D先的DataLoader, 能够永久持续的采样数据loader = torch.utils.data.DataLoader if image_weights else InfiniteDataLoader# Use torch.utils.data.DataLoader() if dataset.properties will update during training else InfiniteDataLoader()dataloader = loader(dataset,batch_size=batch_size,num_workers=nw,sampler=sampler,pin_memory=True,collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn)return dataloader, dataset
\qquad这个函数会在train.py中被调用，用于生成Trainloader, dataset，testloader：
3、自定义DataLoader
\qquad当image_weights=False时（不根据图片样本真实框分布权重来选择图片）就会调用这两个函数 进行自定义DataLoader，进行持续性采样。在上面的create_dataloader模块中被调用。
class InfiniteDataLoader(torch.utils.data.dataloader.DataLoader):""" Dataloader that reuses workers当image_weights=False时就会调用这两个函数 进行自定义DataLoaderhttps://github.com/ultralytics/yolov5/pull/876使用InfiniteDataLoader和_RepeatSampler来对DataLoader进行封装, 代替原先的DataLoader, 能够永久持续的采样数据Uses same syntax as vanilla DataLoader"""def __init__(self, *args, **kwargs):super().__init__(*args, **kwargs)# 调用_RepeatSampler进行持续采样object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))self.iterator = super().__iter__()def __len__(self):return len(self.batch_sampler.sampler)def __iter__(self):for i in range(len(self)):yield next(self.iterator)
class _RepeatSampler(object):""" Sampler that repeats forever这部分是进行持续采样Args:sampler (Sampler)"""def __init__(self, sampler):self.sampler = samplerdef __iter__(self):while True:yield from iter(self.sampler)
4、LoadImagesAndLabels
\qquad这个部分是数据载入（数据增强）部分，也就是自定义数据集部分，继承自Dataset，需要重写__init__,__getitem()__等抽象方法，另外目标检测一般还需要重写collate_fn函数。所以，理解这三个函数是理解数据增强（数据载入）的重中之重。
4.1、init
这个函数的入口是上面的create_dataloader函数：
\qquad其实初始化过程并没有什么实质性的操作，更多是一个定义参数的过程（self参数），以便在__getitem()__中进行数据增强操作，所以这部分代码只需要抓住self中的各个变量的含义就算差不多了。
重点掌握以下红色部分代表什么意思
self.img_files: {list: N} 存放着整个数据集图片的相对路径
self.label_files: {list: N} 存放着整个数据集图片的相对路径
self.labels: 所有图片的所有gt框的信息
self.shapes: 所有图片的shape
self.segments: 所有图片的所有的多边形gt信息
self.batch: 记载着每张图片属于哪个batch
self.n: 数据集中所有图片的数量
self.indices: 记载着所有图片的index
self.rect=True时self.batch_shapes记载每个batch的shape(同一个batch的图片shape相同)，在矩形训练时有用
__init__主要干了一下几件事：
- 赋值一些基础的self变量 用于后面在__getitem__中调用
- 得到path路径下的所有图片的路径self.img_files
- 根据imgs路径找到labels的路径self.label_files
- cache label
- Read cache 生成self.labels、self.shapes、self.img_files、self.label_files、self.batch、self.n、self.indices等变量
- 为Rectangular Training作准备: 生成self.batch_shapes
- 是否需要cache image(一般不需要，太大了)
__init__函数代码：
class LoadImagesAndLabels(Dataset):# for training/testingdef __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False,image_weights=False, cache_images=False, single_cls=False, stride=32, pad=0.0, prefix=''):"""初始化过程并没有什么实质性的操作,更多是一个定义参数的过程（self参数）,以便在__getitem()__中进行数据增强操作,所以这部分代码只需要抓住self中的各个变量的含义就算差不多了self.img_files: {list: N} 存放着整个数据集图片的相对路径self.label_files: {list: N} 存放着整个数据集图片的相对路径cache label -> verify_image_labelself.labels: 如果数据集所有图片中没有一个多边形label labels存储的label就都是原始label(都是正常的矩形label)否则将所有图片正常gt的label存入labels 不正常gt(存在一个多边形)经过segments2boxes转换为正常的矩形labelself.shapes: 所有图片的shapeself.segments: 如果数据集所有图片中没有一个多边形label self.segments=None否则存储数据集中所有存在多边形gt的图片的所有原始label(肯定有多边形label 也可能有矩形正常label 未知数)self.batch: 记载着每张图片属于哪个batchself.n: 数据集中所有图片的数量self.indices: 记载着所有图片的indexself.rect=True时self.batch_shapes记载每个batch的shape(同一个batch的图片shape相同)"""# 1、赋值一些基础的self变量 用于后面在__getitem__中调用self.img_size = img_size # 经过数据增强后的数据图片的大小self.augment = augment # 是否启动数据增强 一般训练时打开 验证时关闭self.hyp = hyp # 超参列表# 图片按权重采样 True就可以根据类别频率(频率高的权重小,反正大)来进行采样 默认False: 不作类别区分self.image_weights = image_weightsself.rect = False if image_weights else rect # 是否启动矩形训练 一般训练时关闭 验证时打开 可以加速self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training)# mosaic增强的边界值 [-320, -320]self.mosaic_border = [-img_size // 2, -img_size // 2]self.stride = stride # 最大下采样率 32self.path = path # 图片路径# 2、得到path路径下的所有图片的路径self.img_files 这里需要自己debug一下 不会太难try:f = [] # image filesfor p in path if isinstance(path, list) else [path]:# 获取数据集路径path，包含图片路径的txt文件或者包含图片的文件夹路径# 使用pathlib.Path生成与操作系统无关的路径，因为不同操作系统路径的‘/’会有所不同p = Path(p) # os-agnostic# 如果路径path为包含图片的文件夹路径if p.is_dir(): # dir# glob.glab: 返回所有匹配的文件路径列表 递归获取p路径下所有文件f += glob.glob(str(p / '**' / '*.*'), recursive=True)# f = list(p.rglob('**/*.*')) # pathlib# 如果路径path为包含图片路径的txt文件elif p.is_file(): # filewith open(p, 'r') as t:t = t.read().strip().splitlines() # 获取图片路径，更换相对路径# 获取数据集路径的上级父目录 os.sep为路径里的分隔符（不同路径的分隔符不同，os.sep可以根据系统自适应）parent = str(p.parent) + os.sepf += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path# f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)else:raise Exception(f'{prefix}{p} does not exist')# 破折号替换为os.sep，os.path.splitext(x)将文件名与扩展名分开并返回一个列表# 筛选f中所有的图片文件self.img_files = sorted([x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in img_formats])# self.img_files = sorted([x for x in f if x.suffix[1:].lower() in img_formats]) # pathlibassert self.img_files, f'{prefix}No images found'except Exception as e:raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {help_url}')# 3、根据imgs路径找到labels的路径self.label_filesself.label_files = img2label_paths(self.img_files) # labels# 4、cache label 下次运行这个脚本的时候直接从cache中取label而不是去文件中取label 速度更快cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache') # cached labels path# Check cacheif cache_path.is_file():# 如果有cache文件，直接加载 exists=True: 是否已从cache文件中读出了nf, nm, ne, nc, n等信息cache, exists = torch.load(cache_path), True # load# 如果图片版本信息或者文件列表的hash值对不上号 说明本地数据集图片和label可能发生了变化 就重新cache label文件if cache.get('version') != 0.3 or cache.get('hash') != get_hash(self.label_files + self.img_files):cache, exists = self.cache_labels(cache_path, prefix), False # re-cacheelse:# 否则调用cache_labels缓存标签及标签相关信息cache, exists = self.cache_labels(cache_path, prefix), False # cache# 打印cache的结果 nf nm ne nc n = 找到的标签数量，漏掉的标签数量，空的标签数量，损坏的标签数量，总的标签数量nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupted, total# 如果已经从cache文件读出了nf nm ne nc n等信息，直接显示标签信息 msgs信息等if exists:d = f"Scanning '{cache_path}' images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupted"tqdm(None, desc=prefix + d, total=n, initial=n) # display all cache resultsif cache['msgs']:logging.info('\n'.join(cache['msgs'])) # display all warnings msg# 数据集没有标签信息 就发出警告并显示标签label下载地址help_urlassert nf > 0 or not augment, f'{prefix}No labels in {cache_path}. Can not train without labels. See {help_url}'# 5、Read cache 从cache中读出最新变量赋给self 方便给forward中使用# cache中的键值对最初有: cache[img_file]=[l, shape, segments] cache[hash] cache[results] cache[msg] cache[version]# 先从cache中去除cache文件中其他无关键值如:'hash', 'version', 'msgs'等都删除[cache.pop(k) for k in ('hash', 'version', 'msgs')] # remove items# pop掉results、hash、version、msgs后只剩下cache[img_file]=[l, shape, segments]# cache.values(): 取cache中所有值 对应所有l, shape, segments# labels: 如果数据集所有图片中没有一个多边形label labels存储的label就都是原始label(都是正常的矩形label)# 否则将所有图片正常gt的label存入labels 不正常gt(存在一个多边形)经过segments2boxes转换为正常的矩形label# shapes: 所有图片的shape# self.segments: 如果数据集所有图片中没有一个多边形label self.segments=None# 否则存储数据集中所有存在多边形gt的图片的所有原始label(肯定有多边形label 也可能有矩形正常label 未知数)# zip 是因为cache中所有labels、shapes、segments信息都是按每张img分开存储的, zip是将所有图片对应的信息叠在一起labels, shapes, self.segments = zip(*cache.values()) # segments: 都是[]self.labels = list(labels) # labels to listself.shapes = np.array(shapes, dtype=np.float64) # image shapes to float64self.img_files = list(cache.keys()) # 更新所有图片的img_files信息 update img_files from cache resultself.label_files = img2label_paths(cache.keys()) # 更新所有图片的label_files信息(因为img_files信息可能发生了变化)if single_cls:for x in self.labels:x[:, 0] = 0n = len(shapes) # number of imagesbi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch indexnb = bi[-1] + 1 # number of batchesself.batch = bi # batch index of imageself.n = n # number of imagesself.indices = range(n) # 所有图片的index# 6、为Rectangular Training作准备# 这里主要是注意shapes的生成 这一步很重要 因为如果采样矩形训练那么整个batch的形状要一样 就要计算这个符合整个batch的shape# 而且还要对数据集按照高宽比进行排序 这样才能保证同一个batch的图片的形状差不多相同 再选则一个共同的shape代价也比较小if self.rect:# Sort by aspect ratios = self.shapes # whar = s[:, 1] / s[:, 0] # aspect ratioirect = ar.argsort() # 根据高宽比排序self.img_files = [self.img_files[i] for i in irect] # 获取排序后的img_filesself.label_files = [self.label_files[i] for i in irect] # 获取排序后的label_filesself.labels = [self.labels[i] for i in irect] # 获取排序后的labelsself.shapes = s[irect] # 获取排序后的whar = ar[irect] # 获取排序后的aspect ratio# 计算每个batch采用的统一尺度 Set training image shapesshapes = [[1, 1]] * nb # nb: number of batchesfor i in range(nb):ari = ar[bi == i] # bi: batch indexmini, maxi = ari.min(), ari.max() # 获取第i个batch中，最小和最大高宽比# 如果高/宽小于1(w > h)，将w设为img_size（保证原图像尺度不变进行缩放）if maxi < 1:shapes[i] = [maxi, 1] # maxi: h相对指定尺度的比例 1: w相对指定尺度的比例# 如果高/宽大于1(w < h)，将h设置为img_size（保证原图像尺度不变进行缩放）elif mini > 1:shapes[i] = [1, 1 / mini]# 计算每个batch输入网络的shape值(向上设置为32的整数倍)# 要求每个batch_shapes的高宽都是32的整数倍，所以要先除以32，取整再乘以32（不过img_size如果是32倍数这里就没必要了）self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride# 7、是否需要cache image 一般是False 因为RAM会不足 cache label还可以 但是cache image就太大了 所以一般不用# Cache images into memory for faster training (WARNING: large datasets may exceed system RAM)self.imgs = [None] * nif cache_images:gb = 0 # Gigabytes of cached imagesself.img_hw0, self.img_hw = [None] * n, [None] * nresults = ThreadPool(num_threads).imap(lambda x: load_image(*x), zip(repeat(self), range(n)))pbar = tqdm(enumerate(results), total=n)for i, x in pbar:self.imgs[i], self.img_hw0[i], self.img_hw[i] = x # img, hw_original, hw_resized = load_image(self, i)gb += self.imgs[i].nbytespbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB)'pbar.close()
4.2、cache_labels
\qquad这个函数用于加载文件路径中的label信息生成cache文件。cache文件中包括的信息有：im_file, l, shape, segments, hash, results, msgs, version等，具体看代码注释。
def cache_labels(self, path=Path('./labels.cache'), prefix=''):"""用在__init__函数中 cache数据集label加载label信息生成cache文件 Cache dataset labels, check images and read shapes:params path: cache文件保存地址:params prefix: 日志头部信息(彩打高亮部分):return x: cache中保存的字典包括的信息有: x[im_file] = [l, shape, segments]一张图片一个label相对应的保存到x, 最终x会保存所有图片的相对路径、gt框的信息、形状shape、所有的多边形gt信息im_file: 当前这张图片的path相对路径l: 当前这张图片的所有gt框的label信息(不包含segment多边形标签) [gt_num, cls+xywh(normalized)]shape: 当前这张图片的形状 shapesegments: 当前这张图片所有gt的label信息(包含segment多边形标签) [gt_num, xy1...]hash: 当前图片和label文件的hash值 1results: 找到的label个数nf, 丢失label个数nm, 空label个数ne, 破损label个数nc, 总img/label个数len(self.img_files)msgs: 所有数据集的msgs信息version: 当前cache version"""x = {} # 初始化最终cache中保存的字典dict# 初始化number missing, found, empty, corrupt, messages# 初始化整个数据集: 漏掉的标签(label)总数量, 找到的标签(label)总数量, 空的标签(label)总数量, 错误标签(label)总数量, 所有错误信息nm, nf, ne, nc, msgs = 0, 0, 0, 0, []desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels..." # 日志# 多进程调用verify_image_label函数with Pool(num_threads) as pool:# 定义pbar进度条# pool.imap_unordered: 对大量数据遍历多进程计算 返回一个迭代器# 把self.img_files, self.label_files, repeat(prefix) list中的值作为参数依次送入(一次送一个)verify_image_label函数pbar = tqdm(pool.imap_unordered(verify_image_label, zip(self.img_files, self.label_files, repeat(prefix))),desc=desc, total=len(self.img_files))# im_file: 当前这张图片的path相对路径# l: [gt_num, cls+xywh(normalized)]# 如果这张图片没有一个segment多边形标签 l就存储原label(全部是正常矩形标签)# 如果这张图片有一个segment多边形标签 l就存储经过segments2boxes处理好的标签(正常矩形标签不处理 多边形标签转化为矩形标签)# shape: 当前这张图片的形状 shape# segments: 如果这张图片没有一个segment多边形标签 存储None# 如果这张图片有一个segment多边形标签 就把这张图片的所有label存储到segments中(若干个正常gt 若干个多边形标签) [gt_num, xy1...]# nm_f(nm): number missing 当前这张图片的label是否丢失 丢失=1 存在=0# nf_f(nf): number found 当前这张图片的label是否存在 存在=1 丢失=0# ne_f(ne): number empty 当前这张图片的label是否是空的 空的=1 没空=0# nc_f(nc): number corrupt 当前这张图片的label文件是否是破损的 破损的=1 没破损=0# msg: 返回的msg信息 label文件完好=‘’ label文件破损=warning信息for im_file, l, shape, segments, nm_f, nf_f, ne_f, nc_f, msg in pbar:nm += nm_f # 累加总number missing labelnf += nf_f # 累加总number found labelne += ne_f # 累加总number empty labelnc += nc_f # 累加总number corrupt labelif im_file:x[im_file] = [l, shape, segments] # 信息存入字典 key=im_file value=[l, shape, segments]if msg:msgs.append(msg) # 将msg加入总msgpbar.desc = f"{desc}{nf} found, {nm} missing, {ne} empty, {nc} corrupted" # 日志pbar.close() # 关闭进度条# 日志打印所有msg信息if msgs:logging.info('\n'.join(msgs))# 一张label都没找到 日志打印help_url下载地址if nf == 0:logging.info(f'{prefix}WARNING: No labels found in {path}. See {help_url}')x['hash'] = get_hash(self.label_files + self.img_files) # 将当前图片和label文件的hash值存入最终字典distx['results'] = nf, nm, ne, nc, len(self.img_files) # 将nf, nm, ne, nc, len(self.img_files)存入最终字典distx['msgs'] = msgs # 将所有数据集的msgs信息存入最终字典distx['version'] = 0.3 # 将当前cache version存入最终字典disttry:torch.save(x, path) # save cache to pathlogging.info(f'{prefix}New cache created: {path}')except Exception as e:logging.info(f'{prefix}WARNING: Cache directory {path.parent} is not writeable: {e}') # path not writeablereturn x
4.3、len
\qquad这个函数是求数据集图片的数量。
def __len__(self):return len(self.img_files)
4.4.、getitem
\qquad这部分是数据增强函数，一般一次性执行batch_size次。
def __getitem__(self, index):"""这部分是数据增强函数，一般一次性执行batch_size次。训练 数据增强: mosaic(random_perspective) + hsv + 上下左右翻转测试 数据增强: letterbox:return torch.from_numpy(img): 这个index的图片数据(增强后) [3, 640, 640]:return labels_out: 这个index图片的gt label [6, 6] = [gt_num, 0+class+xywh(normalized)]:return self.img_files[index]: 这个index图片的路径地址:return shapes: 这个batch的图片的shapes 测试时(矩形训练)才有 验证时为None for COCO mAP rescaling"""# 这里可以通过三种形式获取要进行数据增强的图片index linear, shuffled, or image_weightsindex = self.indices[index]hyp = self.hyp # 超参 包含众多数据增强超参mosaic = self.mosaic and random.random() < hyp['mosaic']# mosaic增强 对图像进行4张图拼接训练 一般训练时运行# mosaic + MixUpif mosaic:# Load mosaicimg, labels = load_mosaic(self, index)# img, labels = load_mosaic9(self, index)shapes = None# MixUp augmentation# mixup数据增强if random.random() < hyp['mixup']: # hyp['mixup']=0 默认为0则关闭 默认为1则100%打开# *load_mosaic(self, random.randint(0, self.n - 1)) 随机从数据集中任选一张图片和本张图片进行mixup数据增强# img: 两张图片融合之后的图片 numpy (640, 640, 3)# labels: 两张图片融合之后的标签label [M+N, cls+x1y1x2y2]img, labels = mixup(img, labels, *load_mosaic(self, random.randint(0, self.n - 1)))# 测试代码 测试MixUp效果# cv2.imshow("MixUp", img)# cv2.waitKey(0)# cv2.destroyAllWindows()# print(img.shape) # (640, 640, 3)# 否则: 载入图片 + Letterbox (val)else:# Load image# 载入图片 载入图片后还会进行一次resize 将当前图片的最长边缩放到指定的大小(512), 较小边同比例缩放# load image img=(343, 512, 3)=(h, w, c) (h0, w0)=(335, 500) numpy index=4# img: resize后的图片 (h0, w0): 原始图片的hw (h, w): resize后的图片的hw# 这一步是将(335, 500, 3) resize-> (343, 512, 3)img, (h0, w0), (h, w) = load_image(self, index)# 测试代码 测试load_image效果# cv2.imshow("load_image", img)# cv2.waitKey(0)# cv2.destroyAllWindows()# print(img.shape) # (640, 640, 3)# Letterbox# letterbox之前确定这张当前图片letterbox之后的shape 如果不用self.rect矩形训练shape就是self.img_size# 如果使用self.rect矩形训练shape就是当前batch的shape 因为矩形训练的话我们整个batch的shape必须统一(在__init__函数第6节内容)shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape# letterbox 这一步将第一步缩放得到的图片再缩放到当前batch所需要的尺度 (343, 512, 3) pad-> (384, 512, 3)# (矩形推理需要一个batch的所有图片的shape必须相同，而这个shape在init函数中保持在self.batch_shapes中)# 这里没有缩放操作，所以这里的ratio永远都是(1.0, 1.0) pad=(0.0, 20.5)img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling# 图片letterbox之后label的坐标也要相应变化 根据pad调整label坐标 并将归一化的xywh -> 未归一化的xyxylabels = self.labels[index].copy()if labels.size: # normalized xywh to pixel xyxy formatlabels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])# 测试代码 测试letterbox效果# cv2.imshow("letterbox", img)# cv2.waitKey(0)# cv2.destroyAllWindows()# print(img.shape) # (640, 640, 3)if self.augment:# Augment imagespaceif not mosaic:# 不做mosaic的话就要做random_perspective增强 因为mosaic函数内部执行了random_perspective增强# random_perspective增强: 随机对图片进行旋转，平移，缩放，裁剪，透视变换img, labels = random_perspective(img, labels,degrees=hyp['degrees'],translate=hyp['translate'],scale=hyp['scale'],shear=hyp['shear'],perspective=hyp['perspective'])# 色域空间增强Augment colorspaceaugment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])# 测试代码 测试augment_hsv效果# cv2.imshow("augment_hsv", img)# cv2.waitKey(0)# cv2.destroyAllWindows()# print(img.shape) # (640, 640, 3)# Apply cutouts 随机进行cutout增强 0.5的几率使用 这里可以自行测试if random.random() < hyp['cutout']: # hyp['cutout']=0 默认为0则关闭 默认为1则100%打开labels = cutout(img, labels)# 测试代码 测试cutout效果# cv2.imshow("cutout", img)# cv2.waitKey(0)# cv2.destroyAllWindows()# print(img.shape) # (640, 640, 3)nL = len(labels) # number of labelsif nL:# xyxy to xywh normalizedlabels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0])# 平移增强 随机左右翻转 + 随机上下翻转if self.augment:# 随机上下翻转 flip up-downif random.random() < hyp['flipud']:img = np.flipud(img) # np.flipud 将数组在上下方向翻转。if nL:labels[:, 2] = 1 - labels[:, 2] # 1 - y_center label也要映射# 随机左右翻转 flip left-rightif random.random() < hyp['fliplr']:img = np.fliplr(img) # np.fliplr 将数组在左右方向翻转if nL:labels[:, 1] = 1 - labels[:, 1] # 1 - x_center label也要映射# 6个值的tensor 初始化标签框对应的图片序号, 配合下面的collate_fn使用labels_out = torch.zeros((nL, 6))if nL:labels_out[:, 1:] = torch.from_numpy(labels) # numpy to tensor# Convert BGR->RGB HWC->CHWimg = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3 x img_height x img_widthimg = np.ascontiguousarray(img) # img变成内存连续的数据 加快运算return torch.from_numpy(img), labels_out, self.img_files[index], shapes
4.5、collate_fn
\qquad很多人以为写完 init 和 getitem 函数数据增强就做完了，我们在分类任务中的确写完这两个函数就可以了，因为系统中是给我们写好了一个collate_fn函数的，但是在目标检测中我们却需要重写collate_fn函数，下面我会仔细的讲解这样做的原因（代码中注释）。
这个函数会在create_dataloader中生成dataloader时调用：
@staticmethoddef collate_fn(batch):"""这个函数会在create_dataloader中生成dataloader时调用：整理函数 将image和label整合到一起:return torch.stack(img, 0): 如[16, 3, 640, 640] 整个batch的图片:return torch.cat(label, 0): 如[15, 6] [num_target, img_index+class_index+xywh(normalized)] 整个batch的label:return path: 整个batch所有图片的路径:return shapes: (h0, w0), ((h / h0, w / w0), pad) for COCO mAP rescalingpytorch的DataLoader打包一个batch的数据集时要经过此函数进行打包 通过重写此函数实现标签与图片对应的划分，一个batch中哪些标签属于哪一张图片,形如[[0, 6, 0.5, 0.5, 0.26, 0.35],[0, 6, 0.5, 0.5, 0.26, 0.35],[1, 6, 0.5, 0.5, 0.26, 0.35],[2, 6, 0.5, 0.5, 0.26, 0.35],]前两行标签属于第一张图片, 第三行属于第二张。。。"""# img: 一个tuple 由batch_size个tensor组成 整个batch中每个tensor表示一张图片# label: 一个tuple 由batch_size个tensor组成 每个tensor存放一张图片的所有的target信息# label[6, object_num] 6中的第一个数代表一个batch中的第几张图# path: 一个tuple 由4个str组成, 每个str对应一张图片的地址信息img, label, path, shapes = zip(*batch) # transposedfor i, l in enumerate(label):l[:, 0] = i # add target image index for build_targets()# 返回的img=[batch_size, 3, 736, 736]# torch.stack(img, 0): 将batch_size个[3, 736, 736]的矩阵拼成一个[batch_size, 3, 736, 736]# label=[target_sums, 6] 6：表示当前target属于哪一张图+class+x+y+w+h# torch.cat(label, 0): 将[n1,6]、[n2,6]、[n3,6]...拼接成[n1+n2+n3+..., 6]# 这里之所以拼接的方式不同是因为img拼接的时候它的每个部分的形状是相同的，都是[3, 736, 736]# 而我label的每个部分的形状是不一定相同的，每张图的目标个数是不一定相同的（label肯定也希望用stack,更方便,但是不能那样拼）# 如果每张图的目标个数是相同的，那我们就可能不需要重写collate_fn函数了return torch.stack(img, 0), torch.cat(label, 0), path, shapes
注意：这个函数一般是当调用了batch_size次 getitem 函数后才会调用一次这个函数，对batch_size张图片和对应的label进行打包。 强烈建议这里大家debug试试这里return的数据是不是我说的这样定义的。
4.6、collate_fn4
\qquad这里是yolo-v5作者实验性的一个代码 quad-collate function 当train.py的opt参数quad=True 则调用collate_fn4代替collate_fn。 作用：将4张mosaic图片[1, 3, 640, 640]合成一张大的mosaic图片[1, 3, 1280, 1280]。将一个batch的图片每四张处理, 0.5的概率将四张图片拼接到一张大图上训练, 0.5概率直接将某张图片上采样两倍训练。
同样在create_dataloader中生成dataloader时调用：
@staticmethoddef collate_fn4(batch):"""同样在create_dataloader中生成dataloader时调用：这里是yolo-v5作者实验性的一个代码 quad-collate function 当train.py的opt参数quad=True 则调用collate_fn4代替collate_fn作用: 如之前用collate_fn可以返回图片[16, 3, 640, 640] 经过collate_fn4则返回图片[4, 3, 1280, 1280]将4张mosaic图片[1, 3, 640, 640]合成一张大的mosaic图片[1, 3, 1280, 1280]将一个batch的图片每四张处理, 0.5的概率将四张图片拼接到一张大图上训练, 0.5概率直接将某张图片上采样两倍训练"""# img: 整个batch的图片 [16, 3, 640, 640]# label: 整个batch的label标签 [num_target, img_index+class_index+xywh(normalized)]# path: 整个batch所有图片的路径# shapes: (h0, w0), ((h / h0, w / w0), pad) for COCO mAP rescalingimg, label, path, shapes = zip(*batch) # transposedn = len(shapes) // 4 # collate_fn4处理后这个batch中图片的个数img4, label4, path4, shapes4 = [], [], path[:n], shapes[:n] # 初始化ho = torch.tensor([[0., 0, 0, 1, 0, 0]])wo = torch.tensor([[0., 0, 1, 0, 0, 0]])s = torch.tensor([[1, 1, .5, .5, .5, .5]]) # scalefor i in range(n): # zidane torch.zeros(16,3,720,1280) # BCHWi *= 4 # 采样 [0, 4, 8, 16]if random.random() < 0.5:# 随机数小于0.5就直接将某张图片上采样两倍训练im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2., mode='bilinear', align_corners=False)[0].type(img[i].type())l = label[i]else:# 随机数大于0.5就将四张图片(mosaic后的)拼接到一张大图上训练im = torch.cat((torch.cat((img[i], img[i + 1]), 1), torch.cat((img[i + 2], img[i + 3]), 1)), 2)l = torch.cat((label[i], label[i + 1] + ho, label[i + 2] + wo, label[i + 3] + ho + wo), 0) * simg4.append(im)label4.append(l)# 后面返回的部分和collate_fn就差不多了 原因和解释都写在上一个函数了 自己debug看一下吧for i, l in enumerate(label4):l[:, 0] = i # add target image index for build_targets()return torch.stack(img4, 0), torch.cat(label4, 0), path4, shapes4
5、img2label_paths
\qquad这个文件是根据数据集中所有图片的路径找到数据集中所有labels对应的路径。用在LoadImagesAndLabels模块的__init__函数中。
def img2label_paths(img_paths):"""用在LoadImagesAndLabels模块的__init__函数中根据imgs图片的路径找到对应labels的路径Define label paths as a function of image paths:params img_paths: {list: 50} 整个数据集的图片相对路径 例如: '..\\datasets\\VOC\\images\\train2007\\000012.jpg'=> '..\\datasets\\VOC\\labels\\train2007\\000012.jpg'"""# 因为python是跨平台的,在Windows上,文件的路径分隔符是'\',在Linux上是'/'# 为了让代码在不同的平台上都能运行，那么路径应该写'\'还是'/'呢？ os.sep根据你所处的平台, 自动采用相应的分隔符号# sa: '\\images\\' sb: '\\labels\\'sa, sb = os.sep + 'images' + os.sep, os.sep + 'labels' + os.sep # /images/, /labels/ substrings# 把img_paths中所以图片路径中的images替换为labelsreturn [sb.join(x.rsplit(sa, 1)).rsplit('.', 1)[0] + '.txt' for x in img_paths]
6、verify_image_label
\qquad这个函数用于检查每一张图片和每一张label文件是否完好。
\qquad 图片文件: 检查内容、格式、大小、完整性
\qquad label文件: 检查每个gt必须是矩形(每行都得是5个数 class+xywh) + 标签是否全部>=0 + 标签坐标xywh是否归一化 + 标签中是否有重复的坐标
verify_image_label函数代码：
def verify_image_label(args):"""用在cache_labels函数中检测数据集中每张图片和每张laebl是否完好图片文件: 内容、格式、大小、完整性label文件: 每个gt必须是矩形(每行都得是5个数 class+xywh) + 标签是否全部>=0 + 标签坐标xywh是否归一化 + 标签中是否有重复的坐标:params im_file: 数据集中一张图片的path相对路径:params lb_file: 数据集中一张图片的label相对路径:params prefix: 日志头部信息(彩打高亮部分):return im_file: 当前这张图片的path相对路径:return l: [gt_num, cls+xywh(normalized)]如果这张图片没有一个segment多边形标签 l就存储原label(全部是正常矩形标签)如果这张图片有一个segment多边形标签 l就存储经过segments2boxes处理好的标签(正常矩形标签不处理 多边形标签转化为矩形标签):return shape: 当前这张图片的形状 shape:return segments: 如果这张图片没有一个segment多边形标签 存储None如果这张图片有一个segment多边形标签 就把这张图片的所有label存储到segments中(若干个正常gt 若干个多边形标签) [gt_num, xy1...]:return nm: number missing 当前这张图片的label是否丢失 丢失=1 存在=0:return nf: number found 当前这张图片的label是否存在 存在=1 丢失=0:return ne: number empty 当前这张图片的label是否是空的 空的=1 没空=0:return nc: number corrupt 当前这张图片的label文件是否是破损的 破损的=1 没破损=0:return msg: 返回的msg信息 label文件完好=‘’ label文件破损=warning信息"""im_file, lb_file, prefix = argsnm, nf, ne, nc = 0, 0, 0, 0 # number missing, found, empty, corrupt labeltry:# 检查这张图片(内容、格式、大小、完整性) verify imagesim = Image.open(im_file) # 打开图片文件im.verify() # PIL verify 检查图片内容和格式是否正常shape = exif_size(im) # 当前图片的大小 image sizeassert (shape[0] > 9) & (shape[1] > 9), f'image size {shape} <10 pixels' # 图片大小必须大于9个pixelsassert im.format.lower() in img_formats, f'invalid image format {im.format}' # 图片格式必须在img_format中if im.format.lower() in ('jpg', 'jpeg'): # 检查jpg格式文件with open(im_file, 'rb') as f:# f.seek: -2 偏移量 向文件头方向中移动的字节数 2 相对位置 从文件尾开始偏移f.seek(-2, 2)# f.read(): 读取图片文件 指令: \xff\xd9 检测整张图片是否完整 如果不完整就返回corrupted JPEGassert f.read() == b'\xff\xd9', 'corrupted JPEG'# verify labelssegments = [] # 存放这张图所有gt框的信息(包含segments多边形: label某一列数大于8)if os.path.isfile(lb_file): # 如果这个label路径存在nf = 1 # label foundwith open(lb_file, 'r') as f: # 读取label文件# 读取当前label文件的每一行: 每一行都是当前图片的一个gtl = [x.split() for x in f.read().strip().splitlines() if len(x)]# any() 函数用于判断给定的可迭代参数 是否全部为False,则返回 False; 如果有一个为 True,则返回True# 如果当前图片的label文件某一列数大于8, 则认为label是存在segment的polygon点(多边形) 就不是矩阵 则将label信息存入segment中if any([len(x) > 8 for x in l]): # is segment# 当前图片中所有gt框的类别classes = np.array([x[0] for x in l], dtype=np.float32)# 获得这张图中所有gt框的label信息(包含segment多边形标签)# 因为segment标签可以是不同长度，所以这里segments是一个列表 [gt_num, xy1...(normalized)]segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in l]# 获得这张图中所有gt框的label信息(不包含segment多边形标签)# segments(多边形) -> bbox(正方形), 得到新标签 [gt_num, cls+xywh(normalized)]l = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1)l = np.array(l, dtype=np.float32) # l: to float32if len(l):# 判断标签是否有五列assert l.shape[1] == 5, 'labels require 5 columns each'# 判断标签是否全部>=0assert (l >= 0).all(), 'negative labels'# 判断标签坐标x y w h是否归一化assert (l[:, 1:] <= 1).all(), 'non-normalized or out of bounds coordinate labels'# 判断标签中是否有重复的坐标assert np.unique(l, axis=0).shape[0] == l.shape[0], 'duplicate labels'else:ne = 1 # label empty l.shape[0] == 0则为空的标签，ne=1l = np.zeros((0, 5), dtype=np.float32)else:nm = 1 # label missing 不存在标签文件，则nm = 1l = np.zeros((0, 5), dtype=np.float32)return im_file, l, shape, segments, nm, nf, ne, nc, ''except Exception as e:nc = 1msg = f'{prefix}WARNING: Ignoring corrupted image and/or label {im_file}: {e}'return [None, None, None, None, nm, nf, ne, nc, msg]
7、load_image
\qquad这个函数是根据图片index，从self或者从对应图片路径中载入对应index的图片 并将原图中hw中较大者扩展到self.img_size, 较小者同比例扩展。会被用在LoadImagesAndLabels模块的__getitem__函数和load_mosaic模块中载入对应index的图片：
load_image函数代码：
def load_image(self, index):"""用在LoadImagesAndLabels模块的__getitem__函数和load_mosaic模块中从self或者从对应图片路径中载入对应index的图片 并将原图中hw中较大者扩展到self.img_size, 较小者同比例扩展loads 1 image from dataset, returns img, original hw, resized hw:params self: 一般是导入LoadImagesAndLabels中的self:param index: 当前图片的index:return: img: resize后的图片(h0, w0): hw_original 原图的hwimg.shape[:2]: hw_resized resize后的图片hw(hw中较大者扩展到self.img_size, 较小者同比例扩展)"""# 按index从self.imgs中载入当前图片, 但是由于缓存的内容一般会不够, 所以我们一般不会用self.imgs(cache)保存所有的图片img = self.imgs[index]# 图片是空的话, 就从对应文件路径读出这张图片if img is None: # not cached 一般都不会使用cache缓存到self.imgs中path = self.img_files[index] # 图片路径img = cv2.imread(path) # 读出BGR图片 (335, 500, 3) HWCassert img is not None, 'Image Not Found ' + pathh0, w0 = img.shape[:2] # orig img hw# img_size 设置的是预处理后输出的图片尺寸 r=缩放比例r = self.img_size / max(h0, w0) # ratio aspectif r != 1: # if sizes are not equal# cv2.INTER_AREA: 基于区域像素关系的一种重采样或者插值方式.该方法是图像抽取的首选方法, 它可以产生更少的波纹# cv2.INTER_LINEAR: 双线性插值,默认情况下使用该方式进行插值 根据ratio选择不同的插值方式# 将原图中hw中较大者扩展到self.img_size, 较小者同比例扩展img = cv2.resize(img, (int(w0 * r), int(h0 * r)),interpolation=cv2.INTER_AREA if r < 1 and not self.augment else cv2.INTER_LINEAR)return img, (h0, w0), img.shape[:2] # img, hw_original, hw_resizedelse:return self.imgs[index], self.img_hw0[index], self.img_hw[index] # img, hw_original, hw_resized
用在LoadImagesAndLabels模块的__getitem__函数和load_mosaic模块中：
执行效果：
8、augment_hsv
\qquad这个函数是关于图片的色域增强模块，图片并不发生移动，所有不需要改变label，只需要 img 增强即可。
augment_hsv模块代码：
def augment_hsv(img, hgain=0.5, sgain=0.5, vgain=0.5):"""用在LoadImagesAndLabels模块的__getitem__函数hsv色域增强 处理图像hsv，不对label进行任何处理:param img: 待处理图片 BGR [736, 736]:param hgain: h通道色域参数 用于生成新的h通道:param sgain: h通道色域参数 用于生成新的s通道:param vgain: h通道色域参数 用于生成新的v通道:return: 返回hsv增强后的图片 img"""if hgain or sgain or vgain:# 随机取-1到1三个实数，乘以hyp中的hsv三通道的系数 用于生成新的hsv通道r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1 # random gainshue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV)) # 图像的通道拆分 h s vdtype = img.dtype # uint8x = np.arange(0, 256, dtype=r.dtype)lut_hue = ((x * r[0]) % 180).astype(dtype) # 生成新的h通道lut_sat = np.clip(x * r[1], 0, 255).astype(dtype) # 生成新的s通道lut_val = np.clip(x * r[2], 0, 255).astype(dtype) # 生成新的v通道# 图像的通道合并 img_hsv=h+s+v 随机调整hsv之后重新组合hsv通道# cv2.LUT(hue, lut_hue) 通道色域变换 输入变换前通道hue 和变换后通道lut_hueimg_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))# no return needed dst:输出图像cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img) # no return needed hsv->bgr
\qquad还要注意的是这个hsv增强是随机生成各个色域参数的，所以每次增强的效果都是不同的：
第一次：变亮
第二次：变暗
这个函数用在LoadImagesAndLabels模块的__getitem__函数中：
另外，这里涉及到的三个变量来自hyp.yaml超参文件：
9、load_mosaic、load_mosaic9
\qquad这两个函数都是mosaic数据增强，只不过load_mosaic函数是拼接四张图，而load_mosaic9函数是拼接九张图。
9.1、load_mosaic
\qquad这个模块就是很有名的mosaic增强模块，几乎训练的时候都会用它，可以显著的提高小样本的mAP。代码是数据增强里面最难的, 也是最有价值的，mosaic是非常非常有用的数据增强trick, 一定要熟练掌握。
load_mosaic模块代码：
def load_mosaic(self, index):"""用在LoadImagesAndLabels模块的__getitem__函数 进行mosaic数据增强将四张图片拼接在一张马赛克图像中 loads images in a 4-mosaic:param index: 需要获取的图像索引:return: img4: mosaic和随机透视变换后的一张图片 numpy(640, 640, 3)labels4: img4对应的target [M, cls+x1y1x2y2]"""# labels4: 用于存放拼接图像（4张图拼成一张）的label信息(不包含segments多边形)# segments4: 用于存放拼接图像（4张图拼成一张）的label信息(包含segments多边形)labels4, segments4 = [], []s = self.img_size # 一般的图片大小# 随机初始化拼接图像的中心点坐标 [0, s*2]之间随机取2个数作为拼接图像的中心坐标yc, xc = [int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border] # mosaic center x, y# 从dataset中随机寻找额外的三张图像进行拼接 [14, 26, 2, 16] 再随机选三张图片的indexindices = [index] + random.choices(self.indices, k=3) # 3 additional image indices# 遍历四张图像进行拼接 4张不同大小的图像 => 1张[1472, 1472, 3]的图像for i, index in enumerate(indices):# load image 每次拿一张图片 并将这张图片resize到self.size(h,w)img, _, (h, w) = load_image(self, index)# place img in img4if i == 0: # top left 原图[375, 500, 3] load_image->[552, 736, 3] hwc# 创建马赛克图像 [1472, 1472, 3]=[h, w, c]img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles# 计算马赛克图像中的坐标信息(将图像填充到马赛克图像中) w=736 h = 552 马赛克图像：(x1a,y1a)左上角 (x2a,y2a)右下角x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)# 计算截取的图像区域信息(以xc,yc为第一张图像的右下角坐标填充到马赛克图像中，丢弃越界的区域) 图像：(x1b,y1b)左上角 (x2b,y2b)右下角x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h # xmin, ymin, xmax, ymax (small image)elif i == 1: # top right# 计算马赛克图像中的坐标信息(将图像填充到马赛克图像中)x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc# 计算截取的图像区域信息(以xc,yc为第二张图像的左下角坐标填充到马赛克图像中，丢弃越界的区域)x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), helif i == 2: # bottom left# 计算马赛克图像中的坐标信息(将图像填充到马赛克图像中)x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)# 计算截取的图像区域信息(以xc,yc为第三张图像的右上角坐标填充到马赛克图像中，丢弃越界的区域)x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)elif i == 3: # bottom right# 计算马赛克图像中的坐标信息(将图像填充到马赛克图像中)x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)# 计算截取的图像区域信息(以xc,yc为第四张图像的左上角坐标填充到马赛克图像中，丢弃越界的区域)x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)# 将截取的图像区域填充到马赛克图像的相应位置 img4[h, w, c]# 将图像img的【(x1b,y1b)左上角 (x2b,y2b)右下角】区域截取出来填充到马赛克图像的【(x1a,y1a)左上角 (x2a,y2a)右下角】区域img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]# 计算pad(当前图像边界与马赛克边界的距离，越界的情况padw/padh为负值) 用于后面的label映射padw = x1a - x1b # 当前图像与马赛克图像在w维度上相差多少padh = y1a - y1b # 当前图像与马赛克图像在h维度上相差多少# labels: 获取对应拼接图像的所有正常label信息(如果有segments多边形会被转化为矩形label)# segments: 获取对应拼接图像的所有不正常label信息(包含segments多边形也包含正常gt)labels, segments = self.labels[index].copy(), self.segments[index].copy()if labels.size:# normalized xywh normalized to pixel xyxy formatlabels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh)segments = [xyn2xy(x, w, h, padw, padh) for x in segments]labels4.append(labels) # 更新labels4segments4.extend(segments) # 更新segments4# Concat/clip labels4 把labels4（[(2, 5), (1, 5), (3, 5), (1, 5)] => (7, 5)）压缩到一起labels4 = np.concatenate(labels4, 0)# 防止越界 label[:, 1:]中的所有元素的值（位置信息）必须在[0, 2*s]之间,小于0就令其等于0,大于2*s就等于2*s out: 返回for x in (labels4[:, 1:], *segments4):np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()# 测试代码 测试前面的mosaic效果# cv2.imshow("mosaic", img4)# cv2.waitKey(0)# cv2.destroyAllWindows()# print(img4.shape) # (1280, 1280, 3)# 随机偏移标签中心，生成新的标签与原标签结合 replicate# img4, labels4 = replicate(img4, labels4)## # 测试代码 测试replicate效果# cv2.imshow("replicate", img4)# cv2.waitKey(0)# cv2.destroyAllWindows()# print(img4.shape) # (1280, 1280, 3)# Augment# random_perspective Augment 随机透视变换 [1280, 1280, 3] => [640, 640, 3]# 对mosaic整合后的图片进行随机旋转、平移、缩放、裁剪，透视变换，并resize为输入大小img_sizeimg4, labels4 = random_perspective(img4, labels4, segments4,degrees=self.hyp['degrees'],translate=self.hyp['translate'],scale=self.hyp['scale'],shear=self.hyp['shear'],perspective=self.hyp['perspective'],border=self.mosaic_border) # border to remove# 测试代码 测试mosaic + random_perspective随机仿射变换效果# cv2.imshow("random_perspective", img4)# cv2.waitKey(0)# cv2.destroyAllWindows()# print(img4.shape) # (640, 640, 3)return img4, labels4
mosaic算法步骤：
1、在 [img_size x 0.5 : img_size x 1.5] 之间随机选择一个拼接中心的坐标（xc, yc）。需要注意的是这里的img_size是我们需要的图片的大小， 而mosaic初步增强得到的图片的shape应该是2倍的img_size.
2、从 [0, len(label)-1] 之间随机选择3张图片的index, 与传入的图片index共同组成4张照片的集合indices.
-------------------------------------------------------------开始剪切img4---------------------------------------------------------------------
3、for 4张图片：
3.0）、如果是第一张图片，就初始化mosaic图片img4
3.1)、 得到mosaic图片的坐标信息（这个坐标区域是用来填充图像的）：左上角(x1a, y1a), (x2a, y2a)右下角
3.2）、得到截取的图像区域的坐标信息：(x1b,y1b)左上角 (x2b,y2b)右下角
3.3）、将图像img的【(x1b,y1b)左上角 (x2b,y2b)右下角】区域截取出来填充到马赛克图像的【(x1a,y1a)左上角 (x2a,y2a)右下角】 注：这里的填充有三种可能的情况，后面会仔细的讨论。
3.4）、计算当前图像边界与马赛克边界的距离，用于后面的label映射
3.5）、拼接4张图像的labels信息为一张labels4
--------------------------------------------到这里就得到了img4[2 x img_size, 2 x img_size, 3]--------------------------------------
4、Concat labels4
5、clip labels4, 防止越界
-------------------------------------------到这里又得到了labels4(相对img4的)---------------------------------------------------------
6、random_perspective随机透视变换（random_perspective Augment），将img4[2 x img_size, 2 x img_size, 3]=>img4 [img_size, img_size, 3]. 这里我就不仔细的介绍随机透视变换了，下一节会详细介绍的。
--------------------------------------------------到这里就得到了img4[img_size, img_size, 3]--------------------------------------------
7、最后retrun img4[img_size, img_size, 3] 和 labels4(相对img4的)
4张图片进行拼接的时候，通常会出现如下三种情况：
效果显示1：mosaic
shape = (1280, 1280, 3)
效果显示2：mosaic + random_perspective
shape = (640, 640, 3)
9.2、load_mosaic9
\qquad这个模块是作者的实验模块，将九张图片拼接在一张马赛克图像中。总体代码流程和load_mosaic4几乎一样，看懂了load_mosaic4再看这个就很简单了、
load_mosaic9模块代码：
def load_mosaic9(self, index):"""用在LoadImagesAndLabels模块的__getitem__函数 替换mosaic数据增强将九张图片拼接在一张马赛克图像中 loads images in a 9-mosaic:param self::param index: 需要获取的图像索引:return: img9: mosaic和仿射增强后的一张图片labels9: img9对应的target"""# labels9: 用于存放拼接图像（9张图拼成一张）的label信息(不包含segments多边形)# segments9: 用于存放拼接图像（9张图拼成一张）的label信息(包含segments多边形)labels9, segments9 = [], []s = self.img_size # 一般的图片大小(也是最终输出的图片大小)# 从dataset中随机寻找额外的三张图像进行拼接 [14, 26, 2, 16] 再随机选三张图片的indexindices = [index] + random.choices(self.indices, k=8) # 8 additional image indicesfor i, index in enumerate(indices):# Load image 每次拿一张图片 并将这张图片resize到self.size(h,w)img, _, (h, w) = load_image(self, index)# 这里和上面load_mosaic函数的操作类似 就是将取出的img图片嵌到img9中(不是真的嵌入 而是找到对应的位置)# place img in img9if i == 0: # centerimg9 = np.full((s * 3, s * 3, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tilesh0, w0 = h, wc = s, s, s + w, s + h # xmin, ymin, xmax, ymax (base) coordinateselif i == 1: # topc = s, s - h, s + w, selif i == 2: # top rightc = s + wp, s - h, s + wp + w, selif i == 3: # rightc = s + w0, s, s + w0 + w, s + helif i == 4: # bottom rightc = s + w0, s + hp, s + w0 + w, s + hp + helif i == 5: # bottomc = s + w0 - w, s + h0, s + w0, s + h0 + helif i == 6: # bottom leftc = s + w0 - wp - w, s + h0, s + w0 - wp, s + h0 + helif i == 7: # leftc = s - w, s + h0 - h, s, s + h0elif i == 8: # top leftc = s - w, s + h0 - hp - h, s, s + h0 - hppadx, pady = c[:2]x1, y1, x2, y2 = [max(x, 0) for x in c] # allocate coords# 和上面load_mosaic函数的操作类似 找到mosaic9增强后的labels9和segments9labels, segments = self.labels[index].copy(), self.segments[index].copy()if labels.size:labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padx, pady) # normalized xywh to pixel xyxy formatsegments = [xyn2xy(x, w, h, padx, pady) for x in segments]labels9.append(labels)segments9.extend(segments)# 生成对应的img9图片(将对应位置的图片嵌入img9中)img9[y1:y2, x1:x2] = img[y1 - pady:, x1 - padx:] # img9[ymin:ymax, xmin:xmax]hp, wp = h, w # height, width previous# Offsetyc, xc = [int(random.uniform(0, s)) for _ in self.mosaic_border] # mosaic center x, yimg9 = img9[yc:yc + 2 * s, xc:xc + 2 * s]# Concat/clip labelslabels9 = np.concatenate(labels9, 0)labels9[:, [1, 3]] -= xclabels9[:, [2, 4]] -= ycc = np.array([xc, yc]) # centerssegments9 = [x - c for x in segments9]for x in (labels9[:, 1:], *segments9):np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()# img9, labels9 = replicate(img9, labels9) # replicate# Augment 同样进行 随机透视变换img9, labels9 = random_perspective(img9, labels9, segments9,degrees=self.hyp['degrees'],translate=self.hyp['translate'],scale=self.hyp['scale'],shear=self.hyp['shear'],perspective=self.hyp['perspective'],border=self.mosaic_border) # border to removereturn img9, labels9
用法和mosaic一样，直接替换即可：
感兴趣的朋友可以试试，不过用的好像并不是很多，效果没mosaic好。
10、random_perspective
\qquad这个函数是进行随机透视变换，对mosaic整合后的图片进行随机旋转、缩放、平移、裁剪，透视变换，并resize为输入大小img_size。
random_perspective函数代码：
def random_perspective(img, targets=(), segments=(), degrees=10, translate=.1,scale=.1, shear=10, perspective=0.0, border=(0, 0)):"""这个函数会用于load_mosaic中用在mosaic操作之后随机透视变换 对mosaic整合后的图片进行随机旋转、缩放、平移、裁剪，透视变换，并resize为输入大小img_size:params img: mosaic整合后的图片img4 [2*img_size, 2*img_size]如果mosaic后的图片没有一个多边形标签就使用targets, segments为空 如果有一个多边形标签就使用segments, targets不为空:params targets: mosaic整合后图片的所有正常label标签labels4(不正常的会通过segments2boxes将多边形标签转化为正常标签) [N, cls+xyxy]:params segments: mosaic整合后图片的所有不正常label信息(包含segments多边形也包含正常gt) [m, x1y1....]:params degrees: 旋转和缩放矩阵参数:params translate: 平移矩阵参数:params scale: 缩放矩阵参数:params shear: 剪切矩阵参数:params perspective: 透视变换参数:params border: 用于确定最后输出的图片大小 一般等于[-img_size, -img_size] 那么最后输出的图片大小为 [img_size, img_size]:return img: 通过透视变换/仿射变换后的img [img_size, img_size]:return targets: 通过透视变换/仿射变换后的img对应的标签 [n, cls+x1y1x2y2] (通过筛选后的)"""# 设定输出图片的 H W# border=-s // 2 所以最后图片的大小直接减半 [img_size, img_size, 3]height = img.shape[0] + border[0] * 2 # # 最终输出图像的Hwidth = img.shape[1] + border[1] * 2 # 最终输出图像的W# ============================ 开始变换 =============================# 需要注意的是，其实opencv是实现了仿射变换的, 不过我们要先生成仿射变换矩阵M# Center 设置中心平移矩阵C = np.eye(3)C[0, 2] = -img.shape[1] / 2 # x translation (pixels)C[1, 2] = -img.shape[0] / 2 # y translation (pixels)# Perspective 设置透视变换矩阵P = np.eye(3)P[2, 0] = random.uniform(-perspective, perspective) # x perspective (about y)P[2, 1] = random.uniform(-perspective, perspective) # y perspective (about x)# Rotation and Scale 设置旋转和缩放矩阵R = np.eye(3) # 初始化R = [[1,0,0], [0,1,0], [0,0,1]] (3, 3)# a: 随机生成旋转角度 范围在(-degrees, degrees)# a += random.choice([-180, -90, 0, 90]) # add 90deg rotations to small rotationsa = random.uniform(-degrees, degrees)# a += random.choice([-180, -90, 0, 90]) # add 90deg rotations to small rotations# s: 随机生成旋转后图像的缩放比例 范围在(1 - scale, 1 + scale)# s = 2 ** random.uniform(-scale, scale)s = random.uniform(1 - scale, 1 + scale)# s = 2 ** random.uniform(-scale, scale)# cv2.getRotationMatrix2D: 二维旋转缩放函数# 参数 angle:旋转角度 center: 旋转中心(默认就是图像的中心) scale: 旋转后图像的缩放比例R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s)# Shear 设置剪切矩阵S = np.eye(3) # 初始化T = [[1,0,0], [0,1,0], [0,0,1]]S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # x shear (deg)S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # y shear (deg)# Translation 设置平移矩阵T = np.eye(3) # 初始化T = [[1,0,0], [0,1,0], [0,0,1]] (3, 3)T[0, 2] = random.uniform(0.5 - translate, 0.5 + translate) * width # x translation (pixels)T[1, 2] = random.uniform(0.5 - translate, 0.5 + translate) * height # y translation (pixels)# Combined rotation matrix @ 表示矩阵乘法 生成仿射变换矩阵MM = T @ S @ R @ P @ C # order of operations (right to left) is IMPORTANT# 将仿射变换矩阵M作用在图片上if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any(): # image changedif perspective:# 透视变换函数 实现旋转平移缩放变换后的平行线不再平行# 参数和下面warpAffine类似img = cv2.warpPerspective(img, M, dsize=(width, height), borderValue=(114, 114, 114))else:# 仿射变换函数 实现旋转平移缩放变换后的平行线依旧平行# image changed img [1472, 1472, 3] => [736, 736, 3]# cv2.warpAffine: opencv实现的仿射变换函数# 参数： img: 需要变化的图像 M: 变换矩阵 dsize: 输出图像的大小 flags: 插值方法的组合（int 类型！）# borderValue: （重点！）边界填充值 默认情况下，它为0。img = cv2.warpAffine(img, M[:2], dsize=(width, height), borderValue=(114, 114, 114))# Visualize 可视化# import matplotlib.pyplot as plt# ax = plt.subplots(1, 2, figsize=(12, 6))[1].ravel()# ax[0].imshow(img[:, :, ::-1]) # base# ax[1].imshow(img2[:, :, ::-1]) # warped# Transform label coordinates# 同样需要调整标签信息n = len(targets)if n:# 判断是否可以使用segment标签: 只有segments不为空时即数据集中有多边形gt也有正常gt时才能使用segment标签 use_segments=True# 否则如果只有正常gt时segments为空 use_segments=Falseuse_segments = any(x.any() for x in segments)new = np.zeros((n, 4)) # [n, 0+0+0+0]# 如果使用的是segments标签(标签中含有多边形gt)if use_segments: # warp segments# 先对segment标签进行重采样# 比如说segment坐标只有100个，通过interp函数将其采样为n个(默认1000)# [n, x1y2...x99y100] 扩增坐标-> [n, 500, 2]# 由于有旋转，透视变换等操作，所以需要对多边形所有角点都进行变换segments = resample_segments(segments)for i, segment in enumerate(segments): # segment: [500, 2] 多边形的500个点坐标xyxy = np.ones((len(segment), 3)) # [1, 1+1+1]xy[:, :2] = segment # [500, 2]# 对该标签多边形的所有顶点坐标进行透视/仿射变换xy = xy @ M.T # transformxy = xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2] # perspective rescale or affine# 根据segment的坐标，取xy坐标的最大最小值，得到边框的坐标 clipnew[i] = segment2box(xy, width, height) # xy [500, 2]# 不使用segments标签 使用正常的矩形的标签targetselse: # warp boxes# 直接对box透视/仿射变换# 由于有旋转，透视变换等操作，所以需要对四个角点都进行变换xy = np.ones((n * 4, 3))xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1xy = xy @ M.T # transform 每个角点的坐标xy = (xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]).reshape(n, 8) # perspective rescale or affine# create new boxesx = xy[:, [0, 2, 4, 6]]y = xy[:, [1, 3, 5, 7]]new = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T# clip 去除太小的target(target大部分跑到图外去了)new[:, [0, 2]] = new[:, [0, 2]].clip(0, width)new[:, [1, 3]] = new[:, [1, 3]].clip(0, height)# filter candidates 过滤target 筛选box# 长和宽必须大于wh_thr个像素 裁剪过小的框(面积小于裁剪前的area_thr) 长宽比范围在(1/ar_thr, ar_thr)之间的限制# 筛选结果 [n] 全是True或False 使用比如: box1[i]即可得到i中所有等于True的矩形框 False的矩形框全部删除i = box_candidates(box1=targets[:, 1:5].T * s, box2=new.T, area_thr=0.01 if use_segments else 0.10)# 得到所有满足条件的targetstargets = targets[i]targets[:, 1:5] = new[i]return img, targets
这个函数会用于load_mosaic中用在mosaic操作之后进行透视变换/仿射变换：
这个函数的参数来自hyp中的5个参数：
效果显示1：mosaic
shape = (1280, 1280, 3)
效果显示2：mosaic + random_perspective
shape = (640, 640, 3)
11、box_candidates
\qquad这个函数用在random_perspective中，是对透视变换后的图片label进行筛选，去除被裁剪过小的框(面积小于裁剪前的area_thr) 还有长和宽必须大于wh_thr个像素，且长宽比范围在(1/ar_thr, ar_thr)之间的限制。
box_candidates模块代码：
def box_candidates(box1, box2, wh_thr=2, ar_thr=20, area_thr=0.1, eps=1e-16):"""用在random_perspective中 对透视变换后的图片label进行筛选去除被裁剪过小的框(面积小于裁剪前的area_thr) 还有长和宽必须大于wh_thr个像素，且长宽比范围在(1/ar_thr, ar_thr)之间的限制Compute candidate boxes: box1 before augment, box2 after augment, wh_thr (pixels), aspect_ratio_thr, area_ratio:params box1: [4, n]:params box2: [4, n]:params wh_thr: 筛选条件 宽高阈值:params ar_thr: 筛选条件 宽高比、高宽比最大值阈值:params area_thr: 筛选条件 面积阈值:params eps: 1e-16 接近0的数 防止分母为0:return i: 筛选结果 [n] 全是True或False 使用比如: box1[i]即可得到i中所有等于True的矩形框 False的矩形框全部删除"""w1, h1 = box1[2] - box1[0], box1[3] - box1[1] # 求出所有box1矩形框的宽和高 [n] [n]w2, h2 = box2[2] - box2[0], box2[3] - box2[1] # 求出所有box2矩形框的宽和高 [n] [n]ar = np.maximum(w2 / (h2 + eps), h2 / (w2 + eps)) # 求出所有box2矩形框的宽高比和高宽比的较大者 [n, 1]# 筛选条件: 增强后w、h要大于2 增强后图像与增强前图像面积比值大于area_thr 宽高比大于ar_thrreturn (w2 > wh_thr) & (h2 > wh_thr) & (w2 * h2 / (w1 * h1 + eps) > area_thr) & (ar < ar_thr) # candidates
12、replicate
\qquad这个函数是随机偏移标签中心，生成新的标签与原标签结合。可以用在load_mosaic里在mosaic操作之后 random_perspective操作之前， 作者默认是关闭的， 自己可以实验一下效果。
replicate模块代码：
def replicate(img, labels):"""可以用在load_mosaic里在mosaic操作之后 random_perspective操作之前 作者默认是关闭的 自己可以实验一下效果随机偏移标签中心，生成新的标签与原标签结合 Replicate labels:params img: img4 因为是用在mosaic操作之后 所以size=[2*img_size, 2*img_size]:params labels: mosaic整合后图片的所有正常label标签labels4(不正常的会通过segments2boxes将多边形标签转化为正常标签) [N, cls+xyxy]:return img: img4 size=[2*img_size, 2*img_size] 不过图片中多了一半的较小gt个数:params labels: labels4 不过另外增加了一半的较小label [3/2N, cls+xyxy]"""h, w = img.shape[:2] # 得到图片的高和宽boxes = labels[:, 1:].astype(int) # 得到所有gt框的矩形坐标 xyxy [N, xyxy]x1, y1, x2, y2 = boxes.T # 左上角: x1 y1 右下角: x2 y2 [N]s = ((x2 - x1) + (y2 - y1)) / 2 # side length (pixels) [N] 得到N个gt的 (w+h)/2 用来衡量gt框的大小# 生成原标签个数一半的新标签 s.size返回ndarray的元素数量for i in s.argsort()[:round(s.size * 0.5)]: # 返回较小(s较小)的一半gt框的index信息x1b, y1b, x2b, y2b = boxes[i] # 得到这一般较小gt框的坐标信息 左上角x1b y1b 右下角x2b y2bbh, bw = y2b - y1b, x2b - x1b # 得到这一般较小gt框的高宽信息# 随机偏移标签中心点 y范围在[0, 图片高-gt框高] x范围在[0, 图片宽-gt框宽]yc, xc = int(random.uniform(0, h - bh)), int(random.uniform(0, w - bw)) # offset x, y# 重新生成这一半的gt框坐标信息(偏移后)x1a, y1a, x2a, y2a = [xc, yc, xc + bw, yc + bh]# 将图片中真实的gt框偏移到对应生成的坐标(一半较小的偏移 较大的不偏移)img[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]# append 原来的labels标签 + 偏移了的标签labels = np.append(labels, [[labels[i, 0], x1a, y1a, x2a, y2a]], axis=0)return img, labels
会用在load_mosaicload_mosaic里在mosaic操作之后 random_perspective操作之前（一般会关闭 具体还要看个人实验）：
执行效果
13、letterbox
letterbox 的img转换部分
\qquad此时：auto=False（需要pad）, scale_fill=False, scale_up=False。
\qquad显然，这部分需要缩放，因为在这之前的load_image部分已经缩放过了（最长边等于指定大小，较短边等比例缩放），那么在letterbox只需要计算出较小边需要填充的pad, 再将较小边两边pad到相应大小（每个batch需要每张图片的大小，这个大小是不相同的）即可。
也可以结合我画的流程图来理解下面的letterbox代码：
letterbox模块代码：
def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):"""用在LoadImagesAndLabels模块的__getitem__函数 只在val时才会使用将图片缩放调整到指定大小Resize and pad image while meeting stride-multiple constraintshttps://github.com/ultralytics/yolov3/issues/232:param img: 原图 hwc:param new_shape: 缩放后的最长边大小:param color: pad的颜色:param auto: True 保证缩放后的图片保持原图的比例 即 将原图最长边缩放到指定大小，再将原图较短边按原图比例缩放（不会失真）False 将原图最长边缩放到指定大小，再将原图较短边按原图比例缩放,最后将较短边两边pad操作缩放到最长边大小（不会失真）:param scale_fill: True 简单粗暴的将原图resize到指定的大小 相当于就是resize 没有pad操作（失真）:param scale_up: True 对于小于new_shape的原图进行缩放,大于的不变False 对于大于new_shape的原图进行缩放,小于的不变:return: img: letterbox后的图片 HWCratio: wh ratios(dw, dh): w和h的pad"""shape = img.shape[:2] # 第一层resize后图片大小[h, w] = [343, 512]if isinstance(new_shape, int):new_shape = (new_shape, new_shape) # (512, 512)# scale ratio (new / old) 1.024 new_shape=(384, 512)r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) # r=1# 只进行下采样 因为上采样会让图片模糊# (for better test mAP) scale_up = False 对于大于new_shape（r<1）的原图进行缩放,小于new_shape（r>1）的不变if not scaleup: # only scale down, do not scale up (for better test mAP)r = min(r, 1.0)# Compute paddingratio = r, r # width, height ratios (1, 1)new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) # wh(512, 343) 保证缩放后图像比例不变dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding dw=0 dh=41if auto: # minimum rectangle 保证原图比例不变，将图像最大边缩放到指定大小# 这里的取余操作可以保证padding后的图片是32的整数倍(416x416)，如果是(512x512)可以保证是64的整数倍dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding dw=0 dh=0elif scaleFill: # stretch 简单粗暴的将图片缩放到指定尺寸dw, dh = 0.0, 0.0new_unpad = (new_shape[1], new_shape[0])ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios# 在较小边的两侧进行pad, 而不是在一侧paddw /= 2 # divide padding into 2 sides 将padding分到上下，左右两侧 dw=0dh /= 2 # dh=20.5# shape:[h, w] new_unpad:[w, h]if shape[::-1] != new_unpad: # resize 将原图resize到new_unpad（长边相同，比例相同的新图）img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) # 计算上下两侧的padding # top=20 bottom=21left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) # 计算左右两侧的padding # left=0 right=0# add border/padimg = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border# img: (384, 512, 3) ratio=(1.0,1.0) 这里没有缩放操作 (dw,dh)=(0.0, 20.5)return img, ratio, (dw, dh)
__getitem__中letterbox 的label转换部分
总结下在val时这里主要是做了三件事：
- load_image将图片从文件中加载出来，并resize到相应的尺寸（最长边等于我们需要的尺寸，最短边等比例缩放）；
- letterbox将之前resize后的图片再pad到我们所需要的放到dataloader中（collate_fn函数）的尺寸（矩形训练要求同一个batch中的图片的尺寸必须保持一致）；
- 将label从相对原图尺寸（原文件中图片尺寸）缩放到相对letterbox pad后的图片尺寸。因为前两部分的图片尺寸发生了变化，同样的我们的label也需要发生相应的变化。
执行效果
14、cutout
\qquadcutout数据增强，给图片随机添加随机大小的方块噪声 ，目的是提高泛化能力和鲁棒性。来自论文： https://arxiv.org/abs/1708.04552。
\qquad更多原理细节请看博客：【YOLO v4】【trick 8】Data augmentation: MixUp、Random Erasing、CutOut、CutMix、Mosic。
\qquad具体要不要使用，概率是多少可以自己实验。
cutout模块代码：
def cutout(image, labels):"""用在LoadImagesAndLabels模块中的__getitem__函数进行cutout增强 v5源码作者默认是没用用这个的 感兴趣的可以测试一下cutout数据增强, 给图片随机添加随机大小的方块噪声 目的是提高泛化能力和鲁棒性实现：随机选择一个固定大小的正方形区域，然后采用全0填充就OK了，当然为了避免填充0值对训练的影响，应该要对数据进行中心归一化操作，norm到0。论文: https://arxiv.org/abs/1708.04552:params image: 一张图片 [640, 640, 3] numpy:params labels: 这张图片的标签 [N, 5]=[N, cls+x1y1x2y2]:return labels: 筛选后的这张图片的标签 [M, 5]=[M, cls+x1y1x2y2] M<N筛选: 如果随机生成的噪声和原始的gt框相交区域占gt框太大 就筛出这个gt框label"""h, w = image.shape[:2] # 获取图片高和宽def bbox_ioa(box1, box2):"""用在cutout中计算box1和box2相交面积与box2面积的比例Returns the intersection over box2 area given box1, box2. box1 is 4, box2 is nx4. boxes are x1y1x2y2:params box1: 传入随机生成噪声 box [4] = [x1y1x2y2]:params box2: 传入图片原始的label信息 [n, 4] = [n, x1y1x2y2]:return [n, 1] 返回一个生成的噪声box与n个原始label的相交面积与b原始label的比值"""box2 = box2.transpose()# Get the coordinates of bounding boxesb1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]# 求box1和box2的相交面积inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * \(np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1)).clip(0)# box面积box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + 1e-16# 返回box1和box2相交面积 与 box2面积之比return inter_area / box2_area# 设置cutout添加噪声的scale create random masksscales = [0.5] * 1 + [0.25] * 2 + [0.125] * 4 + [0.0625] * 8 + [0.03125] * 16 # image size fractionfor s in scales:# 随机生成噪声 宽高mask_h = random.randint(1, int(h * s))mask_w = random.randint(1, int(w * s))# 随机生成噪声 boxxmin = max(0, random.randint(0, w) - mask_w // 2)ymin = max(0, random.randint(0, h) - mask_h // 2)xmax = min(w, xmin + mask_w)ymax = min(h, ymin + mask_h)# 添加随机颜色的噪声 apply random color maskimage[ymin:ymax, xmin:xmax] = [random.randint(64, 191) for _ in range(3)]# 返回没有噪声的label return unobscured labelsif len(labels) and s > 0.03:box = np.array([xmin, ymin, xmax, ymax], dtype=np.float32) # 随机生成的噪声box# 计算生成的一个噪声box与这张图片中所有gt的box做计算 inter_area/label_area [n, 1]ioa = bbox_ioa(box, labels[:, 1:5])# remove>60% obscured labels 不能切的太大 ioa < 0.60 保留cutout噪声遮挡小于60%的标签labels = labels[ioa < 0.60]return labels
在LoadImagesAndLabels模块中的__getitem__函数进行cutout增强：
执行效果：
mixup增强由超参hyp[‘mixup’]控制，0则关闭 默认为1则100%打开（自己实验判断）：
15、mixup
\qquad这个函数是进行mixup数据增强：按比例融合两张图片。论文：https://arxiv.org/pdf/1710.09412.pdf。
\qquad更多原理细节请看博客：【YOLO v4】【trick 8】Data augmentation: MixUp、Random Erasing、CutOut、CutMix、Mosic。
\qquad具体要不要使用，概率是多少可以自己实验。
mixup模块代码：
def mixup(im, labels, im2, labels2):"""用在LoadImagesAndLabels模块中的__getitem__函数进行mixup增强mixup数据增强, 按比例融合两张图片 Applies MixUp augmentation论文: https://arxiv.org/pdf/1710.09412.pdf:params im:图片1 numpy (640, 640, 3):params labels:[N, 5]=[N, cls+x1y1x2y2]:params im2:图片2 (640, 640, 3):params labels2:[M, 5]=[M, cls+x1y1x2y2]:return img: 两张图片mixup增强后的图片 (640, 640, 3):return labels: 两张图片mixup增强后的label标签 [M+N, cls+x1y1x2y2]"""# 随机从beta分布中获取比例,range[0, 1]r = np.random.beta(32.0, 32.0) # mixup ratio, alpha=beta=32.0# 按照比例融合两张图片im = (im * r + im2 * (1 - r)).astype(np.uint8)# 将两张图片标签拼接到一起labels = np.concatenate((labels, labels2), 0)return im, labels
在LoadImagesAndLabels模块中的__getitem__函数进行mixup增强：
执行效果：
mixup增强由超参hyp[‘mixup’]控制，0则关闭 默认为1则100%打开（自己实验判断）：
16、LoadImages、LoadStreams、LoadWebcam
\qquadload 文件夹中的图片/视频 + 用到很少 load web网页中的数据。
全部代码：
class LoadImages: # for inference"""在detect.py中使用load 文件夹中的图片/视频定义迭代器 用于detect.py"""def __init__(self, path, img_size=640, stride=32):p = str(Path(path).absolute()) # os-agnostic absolute path# glob.glab: 返回所有匹配的文件路径列表 files: 提取图片所有路径if '*' in p:# 如果p是采样正则化表达式提取图片/视频, 可以使用glob获取文件路径files = sorted(glob.glob(p, recursive=True)) # globelif os.path.isdir(p):# 如果p是一个文件夹，使用glob获取全部文件路径files = sorted(glob.glob(os.path.join(p, '*.*'))) # direlif os.path.isfile(p):# 如果p是文件则直接获取files = [p] # fileselse:raise Exception(f'ERROR: {p} does not exist')# images: 目录下所有图片的图片名 videos: 目录下所有视频的视频名images = [x for x in files if x.split('.')[-1].lower() in img_formats]videos = [x for x in files if x.split('.')[-1].lower() in vid_formats]# 图片与视频数量ni, nv = len(images), len(videos)self.img_size = img_sizeself.stride = stride # 最大的下采样率self.files = images + videos # 整合图片和视频路径到一个列表self.nf = ni + nv # number of filesself.video_flag = [False] * ni + [True] * nv # 是不是videoself.mode = 'image' # 默认是读image模式if any(videos):# 判断有没有video文件 如果包含video文件，则初始化opencv中的视频模块，cap=cv2.VideoCapture等self.new_video(videos[0]) # new videoelse:self.cap = Noneassert self.nf > 0, f'No images or videos found in {p}. ' \f'Supported formats are:\nimages: {img_formats}\nvideos: {vid_formats}'def __iter__(self):"""迭代器"""self.count = 0return selfdef __next__(self):"""与iter一起用？"""if self.count == self.nf: # 数据读完了raise StopIterationpath = self.files[self.count] # 读取当前文件路径if self.video_flag[self.count]: # 判断当前文件是否是视频# Read videoself.mode = 'video'# 获取当前帧画面，ret_val为一个bool变量，直到视频读取完毕之前都为Trueret_val, img0 = self.cap.read()# 如果当前视频读取结束，则读取下一个视频if not ret_val:self.count += 1self.cap.release()# self.count == self.nf表示视频已经读取完了if self.count == self.nf: # last videoraise StopIterationelse:path = self.files[self.count]self.new_video(path)ret_val, img0 = self.cap.read()self.frame += 1 # 当前读取视频的帧数print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: ', end='')else:# Read imageself.count += 1img0 = cv2.imread(path) # BGRassert img0 is not None, 'Image Not Found ' + pathprint(f'image {self.count}/{self.nf} {path}: ', end='')# Padded resizeimg = letterbox(img0, self.img_size, stride=self.stride)[0]# Convertimg = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB and HWC to CHWimg = np.ascontiguousarray(img)# 返回路径, resize+pad的图片, 原始图片, 视频对象return path, img, img0, self.capdef new_video(self, path):# 记录帧数self.frame = 0# 初始化视频对象self.cap = cv2.VideoCapture(path)# 得到视频文件中的总帧数self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))def __len__(self):return self.nf # number of filesclass LoadStreams:"""load 文件夹中视频流multiple IP or RTSP cameras定义迭代器 用于detect.py"""def __init__(self, sources='streams.txt', img_size=640, stride=32):self.mode = 'stream' # 初始化mode为imagesself.img_size = img_sizeself.stride = stride # 最大下采样步长# 如果sources为一个保存了多个视频流的文件 获取每一个视频流，保存为一个列表if os.path.isfile(sources):with open(sources, 'r') as f:sources = [x.strip() for x in f.read().strip().splitlines() if len(x.strip())]else:# 反之，只有一个视频流文件就直接保存sources = [sources]n = len(sources) # 视频流个数# 初始化图片 fps 总帧数 线程数self.imgs, self.fps, self.frames, self.threads = [None] * n, [0] * n, [0] * n, [None] * nself.sources = [clean_str(x) for x in sources] # clean source names for later# 遍历每一个视频流for i, s in enumerate(sources): # index, source# Start thread to read frames from video stream# 打印当前视频index/总视频数/视频流地址print(f'{i + 1}/{n}: {s}... ', end='')if 'youtube.com/' in s or 'youtu.be/' in s: # if source is YouTube videocheck_requirements(('pafy', 'youtube_dl'))import pafys = pafy.new(s).getbest(preftype="mp4").url # YouTube URLs = eval(s) if s.isnumeric() else s # i.e. s = '0' local webcam 本地摄像头# s='0'打开本地摄像头，否则打开视频流地址cap = cv2.VideoCapture(s)assert cap.isOpened(), f'Failed to open {s}'# 获取视频的宽和长w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))# 获取视频的帧率self.fps[i] = max(cap.get(cv2.CAP_PROP_FPS) % 100, 0) or 30.0 # 30 FPS fallback# 帧数self.frames[i] = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float('inf') # infinite stream fallback# 读取当前画面_, self.imgs[i] = cap.read() # guarantee first frame# 创建多线程读取视频流，daemon表示主线程结束时子线程也结束self.threads[i] = Thread(target=self.update, args=([i, cap]), daemon=True)print(f" success ({self.frames[i]} frames {w}x{h} at {self.fps[i]:.2f} FPS)")self.threads[i].start()print('') # newline# check for common shapes# 获取进行resize+pad之后的shape，letterbox函数默认(参数auto=True)是按照矩形推理进行填充s = np.stack([letterbox(x, self.img_size, stride=self.stride)[0].shape for x in self.imgs], 0) # shapesself.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equalif not self.rect:print('WARNING: Different stream shapes detected. For optimal performance supply similarly-shaped streams.')def update(self, i, cap):# Read stream `i` frames in daemon threadn, f = 0, self.frames[i]while cap.isOpened() and n < f:n += 1# _, self.imgs[index] = cap.read()cap.grab()# 每4帧读取一次if n % 4: # read every 4th framesuccess, im = cap.retrieve()self.imgs[i] = im if success else self.imgs[i] * 0time.sleep(1 / self.fps[i]) # wait timedef __iter__(self):self.count = -1return selfdef __next__(self):self.count += 1if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord('q'): # q to quitcv2.destroyAllWindows()raise StopIteration# Letterboximg0 = self.imgs.copy()img = [letterbox(x, self.img_size, auto=self.rect, stride=self.stride)[0] for x in img0]# Stack 将读取的图片拼接到一起img = np.stack(img, 0)# Convertimg = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB and BHWC to BCHWimg = np.ascontiguousarray(img)return self.sources, img, img0, Nonedef __len__(self):return 0 # 1E12 frames = 32 streams at 30 FPS for 30 yearsclass LoadWebcam: # for inference"""用到很少 load web网页中的数据"""def __init__(self, pipe='0', img_size=640, stride=32):self.img_size = img_sizeself.stride = strideif pipe.isnumeric():pipe = eval(pipe) # local camera# pipe = 'rtsp://192.168.1.64/1' # IP camera# pipe = 'rtsp://username:password@192.168.1.64/1' # IP camera with login# pipe = 'http://wmccpinetop.axiscam.net/mjpg/video.mjpg' # IP golf cameraself.pipe = pipeself.cap = cv2.VideoCapture(pipe) # video capture objectself.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3) # set buffer sizedef __iter__(self):self.count = -1return selfdef __next__(self):self.count += 1if cv2.waitKey(1) == ord('q'): # q to quitself.cap.release()cv2.destroyAllWindows()raise StopIteration# Read frameif self.pipe == 0: # local cameraret_val, img0 = self.cap.read()img0 = cv2.flip(img0, 1) # flip left-rightelse: # IP cameran = 0while True:n += 1self.cap.grab()if n % 30 == 0: # skip framesret_val, img0 = self.cap.retrieve()if ret_val:break# Printassert ret_val, f'Camera Error {self.pipe}'img_path = 'webcam.jpg'print(f'webcam {self.count}: ', end='')# Padded resizeimg = letterbox(img0, self.img_size, stride=self.stride)[0]# Convertimg = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB and HWC to CHWimg = np.ascontiguousarray(img)return img_path, img, img0, Nonedef __len__(self):return 0
在detect.py中使用：
17、hist_equalize
\qquad这个函数是用于对图片进行直方图均衡化处理，但是在yolov5中并没有用到按这个函数，学习了解下就好，不是重点。
hist_equalize模块代码：
def hist_equalize(img, clahe=True, bgr=False):"""yolov5并没有使用直方图均衡化的增强操作 可以自己试试直方图均衡化增强操作 Equalize histogram on BGR image 'img' with img.shape(n,m,3) and range 0-255:params img: 要进行直方图均衡化的原图:params clahe: 是否要生成自适应均衡化图片 默认True 如果是False就生成全局均衡化图片:params bgr: 传入的img图像是否是bgr图片 默认False:return img: 均衡化之后的图片 大小不变 格式RGB"""# 图片BGR/RGB格式 -> YUV格式yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV if bgr else cv2.COLOR_RGB2YUV)if clahe:# cv2.createCLAHE生成自适应均衡化图像c = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))yuv[:, :, 0] = c.apply(yuv[:, :, 0])else:# 全局均衡化yuv[:, :, 0] = cv2.equalizeHist(yuv[:, :, 0]) # equalize Y channel histogramreturn cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR if bgr else cv2.COLOR_YUV2RGB) # convert YUV image to RGB
自行实验：
if __name__ == '__main__':img1 = cv2.imread("F:\yolo_v5\datasets\coco128\images\\train2017\\000000000036.jpg")img2 = hist_equalize(img1)cv2.imshow("hist_equalize_before", img1)cv2.imshow("hist_equalize_after", img2)cv2.waitKey(0)cv2.destroyAllWindows()
实验结果：
18、create_folder
\qquadcreate_folder函数用于创建一个新的文件夹。会用在下面的flatten_recursive函数中。
create_folder函数代码：
def create_folder(path='./new'):"""用在flatten_recursive函数中创建文件夹 Create folder"""# 如果path存在文件夹，则移除if os.path.exists(path):shutil.rmtree(path) # delete output folder# 再从新新建这个文件夹os.makedirs(path) # make new output folder
19、flatten_recursive
\qquad这个模块是将一个文件路径中的所有文件复制到另一个文件夹中 即将image文件和label文件放到一个新文件夹中。
flatten_recursive模块代码：
def flatten_recursive(path='../../datasets/coco128'):"""没用到 不是很重要 自己有用就用将一个文件路径中的所有文件复制到另一个文件夹中 即将image文件和label文件放到一个新文件夹中Flatten a recursive directory by bringing all files to top level"""new_path = Path(path + '_flat') # '..\datasets\coco128_flat'create_folder(new_path)for file in tqdm(glob.glob(str(Path(path)) + '/**/*.*', recursive=True)):# shutil.copyfile: 复制文件到另一个文件夹中shutil.copyfile(file, new_path / Path(file).name)
自己用：
if __name__ == '__main__':flatten_recursive()
效果：
20、extract_boxes
\qquad这个模块是将目标检测数据集转化为分类数据集 ，集体做法: 把目标检测数据集中的每一个gt拆解开 分类别存储到对应的文件当中。
def extract_boxes(path='../../datasets/coco128'):"""自行使用 生成分类数据集将目标检测数据集转化为分类数据集 集体做法: 把目标检测数据集中的每一个gt拆解开 分类别存储到对应的文件当中Convert detection dataset into classification dataset, with one directory per class使用: from utils.datasets import *; extract_boxes():params path: 数据集地址"""path = Path(path) # images dir 数据集文件目录 默认'..\datasets\coco128'# remove existing path / 'classifier' 文件夹shutil.rmtree(path / 'classifier') if (path / 'classifier').is_dir() else Nonefiles = list(path.rglob('*.*')) # 递归遍历path文件下的'*.*'文件n = len(files) # number of filesfor im_file in tqdm(files, total=n):if im_file.suffix[1:] in img_formats: # 必须得是图片文件# imageim0 = cv2.imread(str(im_file)) # BGRim = im0[..., ::-1] # BGR to RGBh, w = im.shape[:2] # 得到这张图片h w# labels 根据这张图片的路径找到这张图片的label路径lb_file = Path(img2label_paths([str(im_file)])[0])if Path(lb_file).exists():with open(lb_file, 'r') as f:lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # 读取label的各行: 对应各个gt坐标for j, x in enumerate(lb): # 遍历每一个gtc = int(x[0]) # class# 生成新'file_name path\classifier\class_index\image_name'# 如: 'F:\yolo_v5\datasets\coco128\images\train2017\classifier\45\train2017_000000000009_0.jpg'f = (path / 'classifier') / f'{c}' / f'{path.stem}_{im_file.stem}_{j}.jpg' # new filename# f.parent: 'F:\yolo_v5\datasets\coco128\images\train2017\classifier\45'if not f.parent.is_dir():# 每一个类别的第一张照片存进去之前 先创建对应类的文件夹f.parent.mkdir(parents=True)b = x[1:] * [w, h, w, h] # box normalized to 正常大小# b[2:] = b[2:].max() pad: rectangle to squareb[2:] = b[2:] * 1.2 + 3 # padb = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(np.int) # xywh to xyxy# 防止b出界 clip boxes outside of imageb[[0, 2]] = np.clip(b[[0, 2]], 0, w)b[[1, 3]] = np.clip(b[[1, 3]], 0, h)assert cv2.imwrite(str(f), im0[b[1]:b[3], b[0]:b[2]]), f'box failure in {f}'
自行使用：
if __name__ == '__main__':extract_boxes()
生成结果：
分类数据集位置：
按类别划分好：
person类（0类）：
21、autosplit
\qquad这个模块是进行自动划分数据集。当使用自己数据集时，可以用这个模块进行自行划分数据集。
autosplit模块代码：
def autosplit(path='../../datasets/coco128/images', weights=(0.9, 0.1, 0.0), annotated_only=False):"""自行使用 自行划分数据集自动将数据集划分为train/val/test并保存 path/autosplit_*.txt filesUsage: from utils.datasets import *; autosplit():params path: 数据集image位置:params weights: 划分权重 默认分别是(0.9, 0.1, 0.0) 对应(train, val, test):params annotated_only: Only use images with an annotated txt file"""path = Path(path) # images dir# 获取images中所有的图片 image files onlyfiles = sum([list(path.rglob(f"*.{img_ext}")) for img_ext in img_formats], [])n = len(files) # number of filesrandom.seed(0) # 随机数种子# assign each image to a split 根据(train, val, test)权重划分原始图片数据集# indices: [n] 0, 1, 2 分别表示数据集中每一张图片属于哪个数据集 分别对应着(train, val, test)indices = random.choices([0, 1, 2], weights=weights, k=n) #txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] # 3 txt files[(path.parent / x).unlink(missing_ok=True) for x in txt] # remove existing txtprint(f'Autosplitting images from {path}' + ', using *.txt labeled images only' * annotated_only)for i, img in tqdm(zip(indices, files), total=n):if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): # check labelwith open(path.parent / txt[i], 'a') as f:f.write('./' + img.relative_to(path.parent).as_posix() + '\n') # add image to txt file
自行使用：
if __name__ == '__main__':autosplit()
划分结果：
22、dataset_stats
\qquad这个模块是统计数据集的信息返回状态字典。包含: 每个类别的图片数量 + 每个类别的实例数量。
dataset_stats模块代码：
def dataset_stats(path='../data/coco128.yaml', autodownload=False, verbose=False):"""yolov5数据集没有用 自行使用这个模块是统计数据集的信息返回状态字典 包含: 每个类别的图片数量 每个类别的实例数量Return dataset statistics dictionary with images and instances counts per split per classUsage: from utils.datasets import *; dataset_stats('coco128.yaml', verbose=True):params path: 数据集信息 data.yaml:params autodownload: Attempt to download dataset if not found locally:params verbose: print可视化打印:return stats: 统计的数据集信息 详细介绍看后面"""def round_labels(labels):# Update labels to integer class and 6 decimal place floatsreturn [[int(c), *[round(x, 6) for x in points]] for c, *points in labels]with open(check_file(path), encoding='utf-8') as f:data = yaml.safe_load(f) # data dict # coco.yaml数据字典data['path'] = "../../datasets/coco128" # 数据集地址check_dataset(data, autodownload) # 检查数据集是否存在 download dataset if missingnc = data['nc'] # 数据集类别 number of classesstats = {'nc': nc, 'names': data['names']} # statistics dictionaryfor split in 'train', 'val', 'test': # 分train、val、test统计数据集信息if data.get(split) is None:stats[split] = None # i.e. no test setcontinuex = []dataset = LoadImagesAndLabels(data[split], augment=False, rect=True) # load datasetif split == 'train': # 训练集直接读取对应位置的label.cache文件 直接就可以读出数据集信息cache_path = Path(dataset.label_files[0]).parent.with_suffix('.cache') # *.cache path# 统计数据集每个图片label中每个类别gt框的个数# x: {list: img_num} 每个list[class_num] 每个图片的label中每个类别gt框的个数for label in tqdm(dataset.labels, total=dataset.n, desc='Statistics'):x.append(np.bincount(label[:, 0].astype(int), minlength=nc))x = np.array(x) # list to numpy矩阵 [img_num, class_num]# 分别统计train、val、test三个数据集的数据信息# 包括: 'image_stats': 字典dict 图片数量total 没有标签的文件个数unlabelled 数据集每个类别的gt个数[80]# 'instance_stats': 字典dict 数据集中所有图片的所有gt个数total 数据集中每个类别的gt个数[80]# 'labels': 字典dict key=数据集中每张图片的文件名 value=每张图片对应的label信息 [n, cls+xywh]stats[split] = {'instance_stats': {'total': int(x.sum()), 'per_class': x.sum(0).tolist()},'image_stats': {'total': dataset.n, 'unlabelled': int(np.all(x == 0, 1).sum()),'per_class': (x > 0).sum(0).tolist()},'labels': [{str(Path(k).name): round_labels(v.tolist())} for k, v inzip(dataset.img_files, dataset.labels)]}# Save, print and return 统计信息statswith open(cache_path.with_suffix('.json'), 'w') as f:json.dump(stats, f) # save stats *.jsonif verbose: # print可视化print(json.dumps(stats, indent=2, sort_keys=False))# print(yaml.dump([stats], sort_keys=False, default_flow_style=False))return stats
自行使用：
if __name__ == '__main__':dataset_stats()
执行效果：
总结
\qquad这个文件主要进行的是数据增强操作。其中1-7小节是和train.py相连的数据载入+数据增强操作；2-15是具体的数据增强实现函数；16节是数据载入模块，包括 load 文件夹中的图片/视频 + 用到很少 load web网页中的数据；最后一个部分是数据集的扩展功能。着重学习前面俩个部分，后面两个部分学习了解下即可。
–2021.08.28 21.54
【YOLOV5-5.x 源码解读】datasets.py相关推荐
- PTMs：QLoRA技巧之源码解读(qlora.py文件)—解析命令与加载参数→数据预处理→模型训练+评估+推理
PTMs:QLoRA技巧之源码解读(qlora.py文件)-解析命令与加载参数→数据预处理→模型训练+评估+推理 目录 QLoRA技巧之源码解读(qlora.py文件)-解析命令与加载参数→数据预处理 ...
- YOLOV5源码解读（数据集加载和增强）
YOLOV5源码解读系列文章目录 数据集加载和增强 loss计算 前言 此篇为yolov5 3.1 版本,官方地址[https://github.com/ultralytics/yolov5] 看源代 ...
- 利用yolov5训练自己的数据集； yolov5的安装与使用 ； yolov5源码解读
*免责声明: 1\此方法仅提供参考 2\搬了其他博主的操作方法,以贴上路径. 3* 场景一:Anconda环境基本操作 场景二:yolov5的使用 场景三:yolo v5训练自己的数据集 场景四:yo ...
- faster rcnn源码解读（四）之数据类型imdb.py和pascal_voc.py（主要是imdb和roidb数据类型的解说）
转载自:faster rcnn源码解读(四)之数据类型imdb.py和pascal_voc.py(主要是imdb和roidb数据类型的解说) - 野孩子的专栏 - 博客频道 - CSDN.NET ht ...
- faster rcnn源码解读（三）train_faster_rcnn_alt_opt.py
转载自:faster rcnn源码解读(三)train_faster_rcnn_alt_opt.py - 野孩子的专栏 - 博客频道 - CSDN.NET http://blog.csdn.net/u ...
- 【YOLOV5-5.x 源码解读】metrics.py
目录 前言 0.导入需要的包 1.fitness 2.ap_per_class.compute_ap 2.1.ap_per_class 2.2.compute_ap 3.ConfusionMatrix ...
- Ubuntu 16.04下Caffe-SSD的应用（四）——ssd_pascal.py源码解读
前言 caffe-ssd所有的训练时的参数,全部由ssd_pascal.py来定义,之后再去调用相关的脚本和函数,所以想要训练自己的数据,首先要明白ssd_pascal.py各个定义参数的大体意思. ...
- 利用yolov7训练自己的数据集； yolov7的安装与使用 ； yolov7源码解读
*免责声明: 1\此方法仅提供参考 2\搬了其他博主的操作方法,以贴上路径. 3* 场景一:Anconda环境基本操作 场景二:yolov7的使用 场景三:yolov7训练自己的数据集 场景四:实用工 ...
- 【YOLOV5-5.x 源码解读】yolov5s.yaml
目录 前言 1.参数配置 2.anchors配置 3.backbone 4.head 5.整体模型 总结 Reference 前言 源码: YOLOv5源码. 链接: [YOLOV5-5.x 源码讲解 ...
- Bert系列（二）——源码解读之模型主体
本篇文章主要是解读模型主体代码modeling.py.在阅读这篇文章之前希望读者们对bert的相关理论有一定的了解,尤其是transformer的结构原理,网上的资料很多,本文内容对原理部分就不做过多 ...
最新文章
- python爬虫案例_推荐上百个github上Python爬虫案例
- nodejs-Buffer(缓冲区)
- 2、leetcode1001 最大连续1的个数
- u8 和 char如何转化_如何编写高效率稳定的单片机代码
- 用 JMeter 测量性能--测试您的 DB2 数据库
- CSS基础(part9)--CSS背景
- Oracle11完全卸载方法
- Linux Install telnet
- 【微型计算机原理与接口技术】指令概述与标志寄存器
- 工作流实现自定义表单
- 【SPSS】包含多元线性回归、聚类分析、判别分析、主成分、相关系数、非参数秩检验的spss使用方法，含有相关例题，可以解决“数学建模”中数据建模的大部分问题
- 解决IE8 iframe透明的问题
- 计算机专业找工作面试面经总结
- 什么是重载？有什么作用？
- IKBC键盘win键失效
- linux进程的高级管理,Linux进程查看，管理工具
- 全志 android 编译,全志A20启动代码流程分析 ——Android
- CFSSL: 证书管理工具：5:理解CSR文件内容
- 分享：Vue.js新手入门指南-0518-v1.0张雅慧（续）
- iphone 开发之过滤html标签