目录
一、base.BaseDataset
1.__init__类初始化
2.get_img_files根据地址获得图片详细地址
3.get_labels（自定义）获取标签数据
4. update_labels指定类别和单分类设定
5.set_rectangle开启批量矩阵训练
6.cache_images加载图片进程可视化
7.load_image内存加载、cache_image_to_disk磁盘加载
8.build_transforms（自定义）图像增强流水线方法
在前几篇文章中记录了data文件夹下的augment.py数据增强文件，分三部分记录了三种主要的图片增强函数类的使用方法和效果。
图片组合变化博文：https://blog.csdn.net/qq_58718853/article/details/143116380
图片随机变化博文：https://blog.csdn.net/qq_58718853/article/details/143189752
LetterBox形状格式化博文：https://blog.csdn.net/qq_58718853/article/details/143233094
上述以及之后所有实验代码上传至Gitee仓库。（会根据博文进度实时更新）：
Gitee链接：https://gitee.com/machine-bai-xue/yolo-source-code-analysis
如果链接失效，访问404拒绝，可以直接在Gitee码云主页搜索——“机器白学”，所有项目中的YOLO源码实验就是本系列所有实验代码。
本文继续解析data文件夹下的 base.py 文件。这里python文件中的代码只编写了一个类——BaseDataset类，这个类是一个用于管理图像数据的数据集模板。它提供了加载、处理和缓存图像的方法，并且可以处理标签和数据增强等任务。
特别的，这个类是一个规范化的模板。意思是此类作为一个样板，其中具体实现功能函数的逻辑还需要自己编写，此类可以作为一个基类使用，在子类中编写具体的功能实现。这么做的一个好处是，给数据加载类的编写提供了十分便利的灵活性，允许用户根据自己的数据便签格式加载成Yolo接受的数据集。
一、base.BaseDataset
1.__init__类初始化
首先记录该基类初始化参数信息，下面表格归纳了该类的初始化参数配置含义。除img_path为必须输入初始化参数外，其余都是可选输入参数。
| img_path | 表示图像文件所在文件夹的路径，用于指向训练所需的图像资源。str字符串类型输入 |
| imgsz | 表示图像的尺寸大小，默认值为640。训练前将所有图像调整为该尺寸，以确保统一输入。int整数类型输入 |
| cache | 表示是否在训练期间将图像缓存到RAM或磁盘中。默认为False，如果设置为True，可以提高加载速度，但会消耗更多的内存或存储空间。bool类型输入 |
| augment | 表示是否启用数据增强。默认为True，启用后可以对图像进行各种变换（如翻转、裁剪等）以增加数据集的多样性，提高模型的泛化能力。bool类型输入 |
| hyp | 表示数据增强所需的超参数。默认为None，若提供该字典，可根据其内容自定义数据增强的细节。dict字典类型输入 |
| prefix | 表示日志消息中的前缀，用于输出时增加标识符，方便分辨不同日志来源。默认为空字符串‘’ |
| rect | 表示是否使用矩形训练（rectangular training），默认为False。启用后将保持图像的长宽比进行训练，但可能会导致尺寸不一致。bool类型输入 |
| batch_size | 表示训练时的批次大小，默认为None。可根据硬件资源和模型需求调整该值。int整数类型输入 |
| stride | 表示使用矩形训练（rect=True）时模型的步幅大小，默认为32。int整数类型输入 |
| pad | 表示使用矩阵训练（rect=True）时，边缘需要填充的大小，默认为0.0。float浮点数输入 |
| single_cls | 表示是否使用单类别训练。默认为False，若设置为True，则数据集中只有一个类别。bool类型输入 |
| classes | 包含所需类别的列表。默认为None，表示使用数据集中的所有类别。可以通过指定类别列表来过滤不需要的类别。list列表类型输入 |
| fraction | 表示所使用的数据集的比例，默认为1.0，表示使用全部数据。通过减少该值，可以只使用部分数据来进行快速测试或减少训练时间。float浮点数输入 |
继续，此类的运行逻辑实际上是在init中实现的。因此阅读init具体代码可以明白此类加载数据的流程。其中初始化中使用了很多后续定义的函数，以下也按照init中使用的顺序介绍其中函数的使用。
首先是根据图片项目地址，获得所有图片信息的函数——self.get_img_files(self.img_path)
2.get_img_files根据地址获得图片详细地址
总体来说，该函数分为两个大的部分，第一部分是下图中红黄蓝三色部分，其主要实现目的是将 img_path 目录下所有文件遍历出来加到 f 列表中。第二部分是绿色部分，主要功能是根据默认图片数据类型后缀 IMG_FORMATS 在目录下所有文件中筛选出图片文件的地址。
3.get_labels（自定义）获取标签数据
继续根据 init 中运行顺序逻辑，下一个是获取便签labels信息的 get_labels 函数。
这是一个可以灵活自定义功能的函数，需要在以BaseDataset为基类的自定义类下定义。 其labels返回的数据类型必须是字典，且要按以下官方要求的内容格式。
于是自定义 MyDataset 类继承BaseDataset，可以参考 get_img_files 代码遍历文件，只需修改最后筛选后缀为‘txt’即可（使用coco8数据，标签数据使用txt保存。如果是其他自定义数据，如json文件只需在这修改后缀即可）。打印取出的地址结果。
然后就只需根据txt标签文件和jpg图片文件的地址读取相应信息，按官方要求格式编写返回labels即可。具体代码参看上传仓库（yolo源码实验仓库）
4. update_labels指定类别和单分类设定
获取便签labels后，init逻辑继续是更新标签labels——update_labels，这么做的目的是实现只取某些类别构建数据集include_class 和 如果是视为单分类问题（重视对于框的位置预测而非框内目标类别）如何重新赋值类别。
下图展示了其具体代码实现的逻辑。
5.set_rectangle开启批量矩阵训练
这个函数的功能是通过一定的规则计算出每个批量里的图片矩阵形状统一大小。通过这种方式，YOLO模型能够在训练时使用统一的输入尺寸，保证批次内部图像尺寸一致。具体功能和作用实现要结合之后的 get_image_and_label 函数来看。
默认为False。
本次解读源码采用打印其过程变量信息形式进行，该函数首先计算每个图片对应的批次索引——假设一共有8个图片样本，批次数设为4。下图可见前四个样本被分为第0批次，后四个被分到了第1批次，一共是2个批次数据。
接着根据高宽比重新排列图片地址和labels信息
然后根据从新排序每个批次里的图片高宽比得到对应每个批次的高宽比（红色箭头指向）。
最后得到每个批次里的图片形状和对应的批次索引。
6.cache_images加载图片进程可视化
还是继续 init 中的顺序逻辑，下载要加载的是图片数据。
首先看 init 上图红框中的代码，先是初始化了三个列表self.ims、self.im_hw0、self.im_hw，分别用来存储图像数据、原始和调整后的图像高宽。每个列表长度和图片数据样本数量一致。
然后再根据参数self.cache设置来选择使用 ‘ram’（内存缓存）、‘disk’（磁盘缓存）还是None（不缓存）。
特别的，如果使用ram缓存，并且启用了hyp.deterministic（确定性训练），则警告ram缓存会导致非确定性结果，建议使用disk缓存。
下面来看cache_images具体加载图片的代码逻辑。实际上此函数最主要的功能是加载过程中全局信息的可视化和掌控——通过多线程将图片缓存到ram内存或disk磁盘，并实时计算和显示缓存的总大小和进度。具体如下图所示。
7.load_image内存加载、cache_image_to_disk磁盘加载
通过阅读 cache_images 中的代码可知，图片具体加载功能的实现是分为内存和磁盘，并分别写在函数 load_images 和 cache_image_to_disk 中的。
先来看 load_images ，它首先检查图片是否已经缓存到ram，支持.npy格式和一般图片格式的数据输入。同时，其还会对加载图片做二次处理，如果启用矩形训练（rect_mode=True），那么会保持图片高宽比缩放图片，否则将图片拉伸为正方形。
同时，如果还开启了训练数据增强 self.augment=True，此方法后续还会将图片索引添加进 self.buffer 动态缓存区。
由于此处 load_image 进行了图片的操作处理，因此单独可视化一下其处理前后效果，方便理解其操作。
首先打印一下 load 中需要用到的一些参数配置。可以看到初始化默认参数cache=False时，是没有将数据加载到ram内存中的。
选定好图片索引后，让我们先测试不使用矩形训练（rect_mode= False）。可以看到根据指定的索引，第二张图已经被正确加载了。
下面两张对比图，左侧是原图，右侧是加载后的输出图。可以看出加载输出图被强制，从原图的 640x426 大小调整为满足模型输入大小 640x640 的“正方形”形状了。
继续，查看使用rect_mode = True 时的效果对比图。可以看到前后没有尺度变化，都是640x426，这是因为原图的宽度640已经满足模型输入的尺度大小，要保持高宽比不变，因此原图不变。
下面改变模型输入尺度——imgsz大小（输入尺度改为320x320），再看看上述启用矩形和不启用的区别。
首先是不启用矩形（rect_mode=False），显而易见是将其直接拉伸到目标尺度。
640x426——> 320x320
而启用矩形则明显不同 rect_mode = True，是先将宽度拉到目标尺度，高度在同比变化。
640x426——> 320x213
继续讨论 cache_image_to_disk，这个方法就简单了，只是保存npy格式数据到本地。
可以看到运行会保存到和原图相同路径的同名 .npy 文件。
8.build_transforms（自定义）图像增强流水线方法
终于最后，读取更新了标签数据、加载调整了图片数据，最后就是根据需要进行数据增强的预处理操作了——build_transforms。这也是整个YOLO数据集类初始化逻辑的最后一步。
这里和获取标签时一样，也是可以自己灵活构建的，构建时需注意三点：1.输入数据应符合augment.py源码解析时的规范要求；2.应使用Compose（[ ]） 的流水线操作图片形式；3.所有图片增强操作中的超参数从 hyp 中获得（默认是from ultralytics.utils import DEFAULT_CFG）
此处测试使用之前数据增强augment博文中的图像翻转处理——RandomFlip和图像色调饱和度色轮调整——RandomHSV。并将图像转换操作写进自己定义的类MyDataset里面。
此时转换操作里的参数都是默认参数值—— DEFAULT_CFG，如果需要修改参数值，要么修改默认配置 cfg 里的信息；要么在初始化类时传入自定义的超参数，使用 hyp=的形式。可以打印默认参数看一下。
运行图片转换操作时，还要注意输入数据格式，这个参考之前关于数据增强的文章中——3._mosaic"n"：Mosaic的三种组合模式。
经过数据增加随机变化后的图片保存效果图如下。
至此，关于YOLO源码文件data下的 base.py 文件解读完成，其给出了一个YOLO数据集加载流程的模板，并可以自定义加载标签数据的格式，自定义数据增强操作。
