YOLOv8的dataloader，数据增强--源码解析

1、入口

在训练下，找创建dataloader的入口，即 self.get_dataloader(.....)，

源代码文件：ultralytics/engine/trainer.py

self.get_dataloader  属于 基类：  class BaseTrainer: 的成员函数。

执行 self.get_dataloader ()，进入了另一个文件：ultralytics/models/yolo/detect/train.py

这里是 子类：  class DetectionTrainer(BaseTrainer):  重写了  父类的  self.get_dataloader ()

整体流程：get_dataloader() --> build_dataset()  --> build_dataloader()

2、为了方便，从源码中单独提取出使用到的 dataset 和 dataloader 源码。

 从main.py 中开始 执行  build_dataset()

dataset_path = r'D:\\Adataset\\VisDrone2019\\images\\train'
batch_size = 32

3、解析  build_dataset() 

①直接进入  build_yolo_dataset() 

self.args ：输入的超参数

img_path：图片所在根目录

self_data：路径，数据集类别等信息

② 进入 子类 class YOLODataset(BaseDataset)

从  super().__init__(*args, **kwargs)  进入其基类 class BaseDataset(Dataset)

在上面 __init()__最后一行定义了对图片处理的代码：

self.transforms = self.build_transforms(hyp=hyp)

子类 class YOLODataset(BaseDataset)中实现了 def build_transforms(self, hyp=None):

下面进入 v8_transforms（...）

③ build_transforms（）、v8_transforms（）

正式进入定义数据增强的代码，来自 ultralytics/data/augment.py

def v8_transforms(dataset, imgsz, hyp, stretch=False):

里面的 Compose 是一个自定义的 执行数据增强的类。

最终返回给 ransforms 的是 Compose，是一个类，在__call__()里面执行数据增强。

在return Compose([,,,,,]) 后，build_dataset()里面所有内容就执行完毕，一直return返回到main.py中。

4、什么时候发生数据增强

在  class BaseDataset(Dataset) 中的__getitem__()，为进行图像处理的入口，

发生于：d1 = dataset[4] 或 for batch_images, batch_labels in dataloader:

① d1 = dataset[4]

a) 增强前

在完成创建dataset后，执行 d1 = dataset[4]，直接进入 __getitem__()

self.get_image_and_label(index)：返回图片和标签，同时在保持高宽比不变下把图片的最长边调整为640，结果如下：

进入  self.transforms()，上面提到，self.transforms是 类Compose的一个对象，传入参数后执行Compose中的__call()__()

在执行数据增强前，看看输入的图片长什么样，执行下面代码：

import matplotlib.pyplot as plt 
import numpy as np
plt.imshow(data['img'][..., ::-1])
​​​​​​​plt.show()

得到：640X360的图片。

b) 之前后

这里只使用 Mosaic 数据增强，其它的方法注释掉了。

回到  v8_transforms():

输入参数: dataset        imgsz=640        p=1.0

Mosaic是一个类对象，继承自 BaseMixTransform

接着 ： data = t(data)，t 为Mosaic，它没有__call__()，于是执行其父类的__call()__()

接着上图

进入self._mix_transform(labels)

_mosaic4 就是选4张图片的Mosaic, _mosaic9就是选9张图片的Mosaic,

这里执行 _mosaic4

源码在  ultralytics/data/augment.py 168行

执行流程：

创建1280X1280的全灰numpy数组 （即imgsz*2） -->  在其中随机一个中点（xc,yc） 

--> 主图放左上角，随机选的图依次放右上角、左下角、右下角，每放一张图都对原来标签的bboex加上了偏移，保证标签不乱。

结果图，1280X1280，由于只执行一次dataset[4]，所以全是同一张图片的混合：

继续执行：

d2 = dataset[2]        d3 = dataset[3]，包含了其它图片。

② 使用 dataloader

先创建 dataloader,

输入参数：dataset 即上文创建的dataset

进入build_dataloader()，在 ultralytics/data/build.py

实例化了一个对象：InfiniteDataLoader，
来自一个自定义的dataloader子类：class InfiniteDataLoader(dataloader.DataLoader)
dataloader.DataLoader 为底层touch代码，这里不管。

每次对一个batch数据处理。