YOLO26 图像分割 - 训练的数据加载：取图、在线增强与标签缓存

YOLO26 图像分割 - 训练的数据加载：取图、在线增强与标签缓存

flyfish

版本：ultralytics 8.4.150

源码：

ultralytics/data/base.py、

ultralytics/data/dataset.py、

ultralytics/data/utils.py

ultralytics 在数据加载上做了三层缓存设计

一、一次 __getitem__ 的旅程

# base.py ------ 数据集取样本的唯一入口
def __getitem__(self, index):
    return self.transforms(self.get_image_and_label(index))

PyTorch 的 DataLoader 对一个 batch 里的每个索引各调用一次 __getitem__（多个 worker 并行），每次只处理一张图：

load_image()        ← 三级取图：内存 → .npy → 磁盘原图；只做等比缩放、补通道
   ↓
self.ims[i]  ←────── 缩放后的"底图"存入缓存（未增强）
   ↓
get_image_and_label() ← 组装标签字典（base.py）
   ↓
transforms(...)      ← 增强在这里实时做：mosaic / 翻转 / HSV / 掩码栅格化
   ↓
collate_fn           ← 把 N 张拼成 batch（dataset.py）

取图发生在增强之前，self.ims 里存的不是增强后的图；

/255 归一化既不在这里、也不在增强里，而是在进模型前的 preprocess_batch（models/yolo/detect/train.py）。

二、取图 内存 → npy → 磁盘原图

load_image 是数据加载的逻辑在 base.py：

# base.py ------ 三级查找
im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i]
if im is None:  # not cached in RAM
    if fn.exists():  # load npy
        try:
            im = np.load(fn)                    # ② .npy 缓存
            ...
        except Exception:
            ...
            im = imread(f, flags=self.cv2_flag) # ③ 磁盘原图
    else:  # read image
        im = imread(f, flags=self.cv2_flag)     # ③ 磁盘原图（BGR）

第一级：内存 self.ims[i]（最快）

self.ims 在数据集初始化时建好（base.py）：

# base.py
self.ims, self.im_hw0, self.im_hw = [None] * self.ni, [None] * self.ni, [None] * self.ni
self.npy_files = [Path(f).with_suffix(".npy") for f in self.im_files]
self.cache = cache.lower() if isinstance(cache, str) else "ram" if cache is True else None

内存里存什么，取决于 cache 参数：

cache=ram：启动时 cache_images() 把全部图片读进内存（base.py，进度条会显示缓存了多 GB），之后每步直接命中 self.ims[i]，完全不碰磁盘。

默认（cache=False）：一个滑动 buffer ，只缓存最近读过的一小批：

# base.py
self.buffer = []  # buffer size = batch size
self.max_buffer_length = min((self.ni, self.batch_size * 8, 1000)) if self.augment else 0

每次读完图后，把这张图的缩放结果写回 self.ims[i] 并记入 buffer（base.py）；buffer 满了就踢最旧的、清空它对应的 ims 槽位。这个 buffer 是为 mosaic 服务的------mosaic 一次要额外抽 3 张随机图，最近读过的图从内存直接拿，不必反复读盘。

第二级：磁盘 .npy（比读原图快）

cache=disk 时，启动阶段 cache_images_to_disk（base.py）把每张原图解码后以 np.save 存成 .npy：

# base.py
f = self.npy_files[i]
if not f.exists():
    np.save(f.as_posix(), imread(self.im_files[i], flags=self.cv2_flag), allow_pickle=False)

.npy 缓存就像是提前做好的半成品：JPEG 或 PNG 图片在硬盘上是以压缩格式存储的，每次读取时电脑都必须先执行 "解码" 这一步 ------ 把压缩数据还原成一格一格的像素，才能拿来训练，这一步相当耗时；而当你开启 cache=disk 时，ultralytics 会在训练启动时把每一张图片只解码一次，并将解压后的像素数组直接保存成一个独立的 .npy 文件，之后每次取图，它只要发现这个文件存在，就能用 np.load 直接把现成的像素数组拿出来用，完全跳过了重复解码的过程，就像半成品菜热一下就能上桌一样；不过这种便利是有代价的，因为解压后的数据比压缩后的原图大得多，整个数据集多存一份 .npy 副本会占用大量磁盘空间，所以 ultralytics 在开启这种缓存前会先做一个粗略的估算 ------ 它随机抽出 30 张图片，量出平均每张解压后占多少字节，再乘以图片总数，并额外加上 50% 的安全余量，以此推断整个缓存大约需要多少磁盘空间；如果估算结果显示剩余空间放不下，它就会自动放弃做 .npy 缓存，打印一条警告告诉你需要多少空间，然后退回最原始的 "每次现读 JPEG 并现场解码" 的模式 ------ 训练功能完全不受影响，只是取图速度变慢了而已check_cache_disk

第三级：磁盘原图

前两级都没命中，才 cv2.imread 读原始文件（BGR），这是最慢的路径。

取到图之后：只做确定性的缩放

# base.py ------ 只缩放，不增强
h0, w0 = im.shape[:2]
if rect_mode:  # resize long side to imgsz while maintaining aspect ratio
    r = self.imgsz / max(h0, w0)
    if r != 1:
        w, h = (min(math.ceil(w0 * r), self.imgsz), min(math.ceil(h0 * r), self.imgsz))
        im = cv2.resize(im, (w, h), interpolation=cv2.INTER_LINEAR)
...
if im.ndim == 2:
    im = im[..., None]  # 灰度图补通道

然后才写进缓存（base.py）：存的是长边缩放到 imgsz、保持宽高比、补好通道的底图,不是原始分辨率，也不是增强后的图。

一张图片本质上是一个多维数组，数组的最后那个维度就叫 通道：彩色图片有三个通道，分别记录红、绿、蓝三种颜色的强度（在 OpenCV 里顺序是 BGR），所以它的形状是 高 × 宽 × 3；而灰度图片只记录明暗，只有一个通道，形状是 高 × 宽，甚至可能只是二维数组。YOLO 模型的第一层卷积，输入通道数是按 3 设计死的，这就要求训练时送进去的每一张图都必须是统一的三通道格式 ------ 如果一批数据里混着 3 通道和 1 通道的图，维度对不上，程序直接就会报错。因此，凡是遇到灰度图，就必须想办法把它 补 成三通道的形态。

补 的做法有两种，本质都一样：把那个唯一的灰度值复制成三份。默认的读图路径用的是 cv2.IMREAD_COLOR 模式，OpenCV 在解码时就会自动把灰度图片处理成三通道，每个像素的三个通道填入完全相同的值 ------ 比如某个像素的亮度是 128，那么它的 B、G、R 就都是 128，三个数字一模一样，因为 RGB 三值相等恰好就表示灰色，正好保留了灰度图的含义。另一种情况是读进来的数组本身就是二维的（比如从 .npy 缓存里取到的数据），这时代码会执行 im[..., None]，在数组末尾加一个维度，把它从 高 × 宽 变成 高 × 宽 × 1，先保证它是三维的、带通道维度的形状，不至于在后续的缩放、拼接、堆叠操作中因为缺少维度而崩溃。

三、self.ims[i] 里存的为什么是未增强的图

# base.py ------ 组装标签
def get_image_and_label(self, index):
    label = deepcopy(self.labels[index])
    ...
    label["img"], label["ori_shape"], label["resized_shape"] = self.load_image(index)
    ...
    return self.update_labels_info(label)

load_image 返回的图直接进入 transforms（也就是 v8_transforms + Format），增强是对这张图实时做的。缓存的写入点（base.py）在 transforms 之前，所以 self.ims[i] 里永远是最新一次读盘的底图。

操作	是否确定性	能否缓存
长边缩放、补通道	确定性（同一张图结果永远一样）	能 缓存，省读盘 + 省缩放
Mosaic 拼图、Flip、HSV、透视	随机（每次迭代结果不同）	不能, 缓存等于每张图只练一个固定花样，增强失效

随机增强的意义就是同一张图每次看到的样子不一样。如果把某一次增强结果固化进缓存，训练数据就变成了静态副本，等于没做增强。这就是在线（online）的含义：增强发生在训练循环内部、每个迭代实时执行；而缩放/读图这类确定性工作可以离线预缓存。

四、栅格化：多边形坐标 → 像素掩码

4.1 什么是栅格化

标注文件里存的不是图，而是"矢量"------一串坐标点围成的多边形：

# labels/train/001.txt 的一行（segment 格式：类别号 后面跟一圈 x,y 点）
3, 150, 120, 260, 110, 270, 135

模型要学、要输出的是掩码------一张格子图，每个格子（像素）标 0/1。栅格化（rasterization）就是把多边形圈住的区域落成一个一个像素格子：落在里面的填 1，外面的填 0。

4.2 代码：在 Format 里实时完成

掩码生成发生在增强流水线最后一步 Format（augment.py），而不是标注时预生成：

# augment.py 附近 ------ _format_segments 调用结构
if self.overlap_mask:  # default.yaml 默认 True
    masks, sorted_idx = polygons2masks_overlap(imgsz, segments, downsample_ratio=self.mask_ratio)
    masks = masks[None]  # (1, H/4, W/4) 实例 ID 掩码
    sem_masks = cls_tensor[masks[0].long() - 1]  # 类别语义掩码
else:
    masks = polygons2masks(imgsz, segments, color=1, downsample_ratio=self.mask_ratio)

4.3 单实例填充：polygon2mask

# utils.py ------ 单个多边形 → 一张二值掩码
mask = np.zeros(imgsz, dtype=np.uint8)              # 全零画布
polygons = np.asarray(polygons, dtype=np.int32)
polygons = polygons.reshape((polygons.shape[0], -1, 2))
cv2.fillPoly(mask, polygons, color=color)           # 多边形内部填 1
nh, nw = (imgsz[0] // downsample_ratio, imgsz[1] // downsample_ratio)
return cv2.resize(mask, (nw, nh))                   # 下采样到 1/mask_ratio

三步：全零画布 → cv2.fillPoly 把多边形内部涂成 1 → 按 mask_ratio 缩小。默认 mask_ratio=4（default.yaml），所以 640×640 的图生成 160×160 的掩码。

4.4 合并实例：polygons2masks_overlap

一张图通常有多个实例，overlap_mask=True（default.yaml）时把它们合成一张实例 ID 掩码：

# utils.py ------ 多实例 → 实例 ID 掩码（按面积排序）
masks = np.zeros((H // r, W // r), dtype=np.int32 if len(segments) > 255 else np.uint8)
for segment in segments:
    mask = polygon2mask(imgsz, [segment.reshape(-1)], downsample_ratio=r, color=1)
    ms.append(mask)
    areas.append(mask.sum())
index = np.argsort(-areas)                          # 按面积降序
for i in range(len(segments)):
    np.maximum(masks, ms[i] * (i + 1), out=masks)   # 第 i 大实例的像素填 i+1

要点：

• 掩码的像素值 = 实例编号（按面积从大到小 1, 2, 3...），不是 0/1 二值；
• 实例重叠处取 maximum（编号更大的小实例盖住大实例）；
• 注释里专门说明用 np.maximum 而不是 masks + mask：老写法在实例超过 128 个时 uint8 会溢出（2*i+1 越界）；
• 实例数 >255 时自动切 int32。

如果标注时就把掩码栅格化存好，Mosaic 拼图、透视变换、翻转之后，掩码和图像像素就对不上了。所以 ultralytics 的选择是：多边形坐标跟随增强做同样的仿射变换，最后一步才栅格化

没有实例的图（负样本）不进入实例循环，掩码直接全 0：

# augment.py 附近 ------ 空图分支
masks = torch.zeros(1, h // 4, w // 4)   # 全 0 实例 ID 掩码

空图在损失侧只贡献背景 cls，掩码本身是空张量。

五、labels/train.cache 与 val.cache：标签扫描的缓存

它是标签扫描结果的序列化缓存。命名规则在 dataset.py：

# dataset.py ------ 缓存路径 = 标签目录的同级 .cache 文件
cache_path = Path(label_files[0]).parent.with_suffix(".cache")

训练集 → labels/train.cache，验证集 → labels/val.cache。内容是 dict：

# dataset.py ------ 保存时写入 hash 和 version
x["hash"] = self.get_cache_hash()
save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)

完整结构：

字段	含义
version	缓存格式版本号（DATASET_CACHE_VERSION，dataset.py）
hash	数据集指纹
labels	每张图的 im_file、shape、cls、bboxes、segments、keypoints 等
results	扫描统计：(nf 找到, nm 缺标签, ne 空标签, nc 损坏, n 总数)
msgs	扫描警告

保存用 np.save，读取用 np.load(allow_pickle=True)（utils.py），本质是 pickle 化的 numpy 文件。

# base.py ------ 数据集构造时取标签
self.labels = self.get_labels()
# dataset.py ------ 有可用缓存就加载，否则重新扫描
cache, exists = load_dataset_cache_file(cache_path), True
assert cache["version"] == DATASET_CACHE_VERSION
assert cache["hash"] == cache_hash
# 任一断言失败 → self.cache_labels(cache_path) 全量重扫并保存

首次训练：没有 train.cache → cache_labels（dataset.py）用线程池并行执行 verify_image_label 逐张校验图片、解析标签 → 保存缓存，日志输出 New cache created: ...，同时能看到扫描进度条。

之后每次启动：缓存存在且 version、hash 都匹配 → 直接 np.load 秒开，跳过扫描。

# dataset.py
def get_cache_hash(self):
    return get_hash(self.label_files + self.im_files)
# utils.py
def get_hash(paths):
    size = 0
    for p in paths:
        size += os.stat(p).st_size     # 文件大小
    h = hashlib.sha256(str(size).encode())
    h.update("".join(paths).encode())  # 文件路径
    return h.hexdigest()

hash 只由文件路径 + 文件总字节数决定------不看文件内容，不看修改时间。