Bootstrap

YOLOv10改进 | 代码逐行解析(三) | YOLO中的Mosaic增强详解(新手入门必读系列)

一、本文介绍

本文给大家带来的是YOLOv10中的Mosaic增强代码的详解,可能有部分人对于这一部分比较陌生,有的读者可能知道Mosaic增强但是不知道其工作原理,具体来说Mosaic增强就是指我们的数据集中的图片在输入给模型之前的一个处理过程(我们的图片并不是直接就输入给模型了,大家的训练结果中的结果检测图片大家可以看到数据集中多个图片会组合在一起这就是简单的Mosaic增强),下面我就来讲解一下其在YOLOv10中工作原理和代码定义,下面图片为一个Mosaic增强后的图片。 

 专栏回顾:YOLOv10改进系列专栏——本专栏持续复习各种顶会内容——科研必备 


目录

一、本文介绍

二、代码详解

2.1 BaseMixTransform

2.2 Mosaic

三、图片示例 

四、 Mosaic 数据增强的优点和缺点

五、本文总结


二、代码详解

首先在开始之前我们要知道Mosaic增强代码在我们项目文件中存在的位置,我们可以按照如下文件中找到'ultralytics/data/augment.py',在这个文件


2.1 BaseMixTransform

在上面文件中我们可以找到一个父类BaseMixTransform(后面会被其它子类继承 | 其中只实现了部分的简单功能,如果大家不明白父类和子类的关系那么大家需要补一补python的基础,这些和深度学习是没有关系的),

class BaseMixTransform:
    """
    Class for base mix (MixUp/Mosaic) transformations.
    基于混合(MixUp/Mosaic)变换的基类。

    This implementation is from mmyolo.
    该实现来自 mmyolo。
    """

    def __init__(self, dataset, pre_transform=None, p=0.0) -> None:
        """Initializes the BaseMixTransform object with dataset, pre_transform, and probability.
        初始化 BaseMixTransform 对象,包含数据集、预处理变换和概率。
        """
        self.dataset = dataset  # 数据集对象,包含图像和标签。
        self.pre_transform = pre_transform  # 预处理变换函数(YOLO选择的为Mosaic)。
        self.p = p  # 应用 MixUp/Mosaic 变换的概率。

    def __call__(self, labels):
        """Applies pre-processing transforms and mixup/mosaic transforms to labels data.
        应用预处理变换和 MixUp/Mosaic 变换到标签数据。
        """
        if random.uniform(0, 1) > self.p:  # 如果随机数大于 p,则不进行数据增强。
            return labels  # 概率决定是否进行数据增强。

        # Get index of one or three other images
        indexes = self.get_indexes()  # 调用 get_indexes 方法获取需要混合的其他图像的索引。
        if isinstance(indexes, int):  # 如果索引是整数,则转换为列表。
            indexes = [indexes]

        # Get images information will be used for Mosaic or MixUp
        mix_labels = [self.dataset.get_image_and_label(i) for i in indexes]  # 根据索引从数据集中获取相应的图像和标签。

        if self.pre_transform is not None:  # 如果定义了预处理变换 pre_transform,则对获取的图像标签进行预处理(Mosaic预处理)。
            for i, data in enumerate(mix_labels):  # 对每个混合标签进行预处理。
                mix_labels[i] = self.pre_transform(data)
        labels["mix_labels"] = mix_labels  # 将获取的图像标签存储在 labels["mix_labels"] 中,这里的图像标签指的是我们对图像标注的标签。

        # Update cls and texts
        labels = self._update_label_text(labels)  # 调用 _update_label_text 方法更新标签文本。
        # Mosaic or MixUp
        labels = self._mix_transform(labels)  # 调用 _mix_transform 方法进行 MixUp 或 Mosaic 变换,实际是Mosaic。
        labels.pop("mix_labels", None)  # 移除 labels 中的 mix_labels 键。
        return labels  # 返回更新后的标签。

    def _mix_transform(self, labels):
        """Applies MixUp or Mosaic augmentation to the label dictionary.
        对标签字典应用 MixUp 或 Mosaic 增强。
        """
        raise NotImplementedError  # 未实现的方法,在子类中具体实现。

    def get_indexes(self):
        """Gets a list of shuffled indexes for mosaic augmentation.
        获取用于 Mosaic 增强的混合索引列表。
        """
        raise NotImplementedError  # 未实现的方法,在子类中具体实现。

    def _update_label_text(self, labels):
        # 这个代码说白了就是我们将多个图片融合到一起的时候,
        # 每个图片都有多个标签,那么我们就需要删除相同的标签因为一个标签我们有一个就可以了,在识别的时候调用就可以了。
        # 同时因为我们的标签列表更新了,原先的每个图片中类别对应的索引也要改变都是通过以下代码的完成的。
        """Update label text.
        更新标签文本。
        """
        if "texts" not in labels:  # 如果 labels 中不包含 texts 键,直接返回 labels。
            return labels  # 如果没有文本标签,直接返回。

        mix_texts = sum([labels["texts"]] + [x["texts"] for x in labels["mix_labels"]], [])  # 合并原始标签文本和混合标签文本。
        mix_texts = list({tuple(x) for x in mix_texts})  # 去重。
        text2id = {text: i for i, text in enumerate(mix_texts)}  # 创建从文本到索引的映射字典 text2id。

        for label in [labels] + labels["mix_labels"]:  # 遍历原始标签和混合标签。
            for i, cls in enumerate(label["cls"].squeeze(-1).tolist()):  # 更新每个标签中的类别索引。
                text = label["texts"][int(cls)]  # 根据类别索引获取文本。
                label["cls"][i] = text2id[tuple(text)]  # 更新类别索引为新的文本索引。
            label["texts"] = mix_texts  # 更新标签文本。
        return labels  # 返回更新后的标签。

上面的代码中主要重要的就是_update_label_text,我们下面通过一个小例子来理解下:

这个函数_update_label_text的作用是更新标签中的文本信息,当进行了Mosaic 数据增强操作后,需要重新整合和去重文本信息,并更新类别索引。

举例解释:

假设我们有以下原始标签 labels 和混合标签 mix_labels:

labels = {
    "cls": [[0], [1]],  # 类别索引
    "texts": [["cat"], ["dog"]]  # 对应的文本标签
}

mix_labels = [
    {
        "cls": [[0], [2]],  # 混合标签的类别索引
        "texts": [["cat"], ["bird"]]  # 混合标签的文本标签
    }
]

labels["mix_labels"] = mix_labels

在调用 _update_label_text函数前:

labels["cls"] 表示有两类:0("cat")和 1("dog")。
labels["texts"] 对应类别的文本标签。
mix_labels包含另一个标签集,有两类:0("cat")和 2("bird")。

 函数步骤解析

1. 检查 "texts" 键是否存在:

if "texts" not in labels:
    return labels

    如果labels中不包含texts键,直接返回labels。在这个例子中,labels包含 texts键,所以继续执行。

2. 合并原始标签文本和混合标签文本:

mix_texts = sum([labels["texts"]] + [x["texts"] for x in labels["mix_labels"]], [])

labels[texts]mix_labels 中的所有 texts 合并为一个列表: 

mix_texts = [["cat"], ["dog"], ["cat"], ["bird"]]

3. 去重:

mix_texts = list({tuple(x) for x in mix_texts})

将合并后的文本转换为集合以去重,然后再转换回列表:

mix_texts = [["cat"], ["dog"], ["bird"]]

4. 创建从文本到索引的映射字典:

text2id = {text: i for i, text in enumerate(mix_texts)}

    创建一个从文本到新索引的字典:

text2id = {text: i for i, text in enumerate(mix_texts)}

5. 更新标签中的类别索引:

for label in [labels] + labels["mix_labels"]:
    for i, cls in enumerate(label["cls"].squeeze(-1).tolist()):
        text = label["texts"][int(cls)]
        label["cls"][i] = text2id[tuple(text)]
    label["texts"] = mix_texts

遍历labels和 mix_labels,将每个标签中的类别索引更新为新的索引:

    对 labels:

  • 原始cls是 [[0], [1]],对应texts是 [[cat], [dog]]。
  • 更新后cls变为 [[0], [1]](没有变化,因为cat和dog的索引没有变化)。
  • 更新texts为 [[cat], [dog], [bird]].

    对 mix_labels:

  • 原始cls是 [[0], [2]],对应 `texts` 是 [[cat], [bird]]。
  • 更新后cls变为 [[0], [2]](同样保持不变,因为cat和bird的新索引与原始索引相同)。
  • 更新texts为 [[cat], [dog], [bird]].

最终,更新后的labels为:

labels = {
    "cls": [[0], [1]],  # 类别索引
    "texts": [["cat"], ["dog"], ["bird"]],  # 更新后的文本标签
    "mix_labels": [
        {
            "cls": [[0], [2]],  # 更新后的混合标签类别索引
            "texts": [["cat"], ["dog"], ["bird"]]  # 更新后的文本标签
        }
    ]
}

这个函数确保在进行数据增强操作后,标签文本是统一且没有重复的,同时更新了类别索引以反映新的文本索引,这个代码其实只实现了这一个功能,绝大多数功能都是在下一小节Mosaic代码中实现的。


2.2 Mosaic

上面说到的是BaseMixTransform父类,下面讲的是子类Mosaic(这里额外解释一下,我们模型定义中是先调用的子类,然后在调用父类中的功能),这个代码同样在'ultralytics/data/augment.py'文件中。

class Mosaic(BaseMixTransform):
    """
    Mosaic augmentation.
    马赛克增强。

    This class performs mosaic augmentation by combining multiple (4 or 9) images into a single mosaic image.
    这个类通过将多个(4 或 9)图像组合成一个单一的马赛克图像来进行增强。
    The augmentation is applied to a dataset with a given probability.
    该增强操作以一定概率应用于数据集。

    Attributes:
        dataset: The dataset on which the mosaic augmentation is applied.
        数据集:应用马赛克增强的数据集。
        imgsz (int, optional): Image size (height and width) after mosaic pipeline of a single image. Default to 640.
        imgsz(int,可选):马赛克处理后单张图像的大小(高度和宽度)。默认为 640(为我们输入给图像的固定尺寸640 x 640的图像无论我们的图片是多么大和多么小都会被转化为640尺寸)。
        p (float, optional): Probability of applying the mosaic augmentation. Must be in the range 0-1. Default to 1.0.
        p(float,可选):应用马赛克增强的概率。必须在 0-1 之间。默认值为 1.0。
        n (int, optional): The grid size, either 4 (for 2x2) or 9 (for 3x3).
        n(int,可选):网格大小,4(表示 2x2 即四张图片合成一个图片)或 9(表示 3x3 9张图片合成一个图片)。
    """

    def __init__(self, dataset, imgsz=640, p=1.0, n=4):
        """Initializes the object with a dataset, image size, probability, and border.
        用数据集、图像大小、概率和边界初始化对象。
        """
        assert 0 <= p <= 1.0, f"The probability should be in range [0, 1], but got {p}."
        断言概率在 [0, 1] 范围内,但得到的是 {p}。
        assert n in {4, 9}, "grid must be equal to 4 or 9."
        断言网格必须等于 4 或 9。
        super().__init__(dataset=dataset, p=p)
        self.dataset = dataset
        self.imgsz = imgsz
        self.border = (-imgsz // 2, -imgsz // 2)  # width, height 宽度,高度
        self.n = n

    def get_indexes(self, buffer=True):
        """Return a list of random indexes from the dataset.
        返回数据集中的随机索引列表。
        """
        if buffer:  # select images from buffer
            return random.choices(list(self.dataset.buffer), k=self.n - 1)
            # 从缓冲区中选择图像(我们的数据加载过程是定义在CPU上,如果大家学过计组就会知道缓冲区的概念)
        else:  # select any images
            return [random.randint(0, len(self.dataset) - 1) for _ in range(self.n - 1)]
            # 选择任意图像,随机算则图像作为处理图片

    def _mix_transform(self, labels):
        """Apply mixup transformation to the input image and labels.
        对输入图像和标签应用混合变换。
        """
        assert labels.get("rect_shape", None) is None, "rect and mosaic are mutually exclusive."
        断言标签中没有 "rect_shape" 键,"rect" 和 "mosaic" 是互斥的。
        assert len(labels.get("mix_labels", [])), "There are no other images for mosaic augment."
        断言有其他图像用于马赛克增强。
        return (
            self._mosaic3(labels) if self.n == 3 else self._mosaic4(labels) if self.n == 4 else self._mosaic9(labels)
        )  # 根据 n 的值选择相应的马赛克方法

    def _mosaic3(self, labels):
        """Create a 1x3 image mosaic.
        创建一个 1x3 的马赛克图像。
        """
        mosaic_labels = []
        s = self.imgsz
        for i in range(3):
            labels_patch = labels if i == 0 else labels["mix_labels"][i - 1]
            # 加载图像
            img = labels_patch["img"]
            h, w = labels_patch.pop("resized_shape")

            # 将图像放置在 img3 中
            if i == 0:  # 中间
                img3 = np.full((s * 3, s * 3, img.shape[2]), 114, dtype=np.uint8)  # 基础图像有 3 块
                h0, w0 = h, w
                c = s, s, s + w, s + h  # xmin, ymin, xmax, ymax (基础坐标)
            elif i == 1:  # 右侧
                c = s + w0, s, s + w0 + w, s + h
            elif i == 2:  # 左侧
                c = s - w, s + h0 - h, s, s + h0

            padw, padh = c[:2]
            x1, y1, x2, y2 = (max(x, 0) for x in c)  # 分配坐标

            img3[y1:y2, x1:x2] = img[y1 - padh :, x1 - padw :]  # img3[ymin:ymax, xmin:xmax]

            # 假设 imgsz*2 的马赛克尺寸
            labels_patch = self._update_labels(labels_patch, padw + self.border[0], padh + self.border[1])
            mosaic_labels.append(labels_patch)
        final_labels = self._cat_labels(mosaic_labels)

        final_labels["img"] = img3[-self.border[0] : self.border[0], -self.border[1] : self.border[1]]
        return final_labels

    def _mosaic4(self, labels):
        """Create a 2x2 image mosaic.
        创建一个 2x2 的马赛克图像。
        """
        mosaic_labels = []
        s = self.imgsz
        yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.border)  # 马赛克中心 x, y
        for i in range(4):
            labels_patch = labels if i == 0 else labels["mix_labels"][i - 1]
            # 加载图像
            img = labels_patch["img"]
            h, w = labels_patch.pop("resized_shape")

            # 将图像放置在 img4 中
            if i == 0:  # 左上角
                img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # 基础图像有 4 块
                x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (大图)
                x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (小图)
            elif i == 1:  # 右上角
                x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
                x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
            elif i == 2:  # 左下角
                x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
                x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
            elif i == 3:  # 右下角
                x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
                x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)

            img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]  # img4[ymin:ymax, xmin:xmax]
            padw = x1a - x1b
            padh = y1a - y1b

            labels_patch = self._update_labels(labels_patch, padw, padh)
            mosaic_labels.append(labels_patch)
        final_labels = self._cat_labels(mosaic_labels)
        final_labels["img"] = img4
        return final_labels

    def _mosaic9(self, labels):
        """Create a 3x3 image mosaic.
        创建一个 3x3 的马赛克图像。
        """
        mosaic_labels = []
        s = self.imgsz
        hp, wp = -1, -1  # 上一次的高度、宽度
        for i in range(9):
            labels_patch = labels if i == 0 else labels["mix_labels"][i - 1]
            # 加载图像
            img = labels_patch["img"]
            h, w = labels_patch.pop("resized_shape")

            # 将图像放置在 img9 中
            if i == 0:  # 中间
                img9 = np.full((s * 3, s * 3, img.shape[2]), 114, dtype=np.uint8)  # 基础图像有 9 块
                h0, w0 = h, w
                c = s, s, s + w, s + h  # xmin, ymin, xmax, ymax (基础坐标)
            elif i == 1:  # 上方
                c = s, s - h, s + w, s
            elif i == 2:  # 右上角
                c = s + wp, s - h, s + wp + w, s
            elif i == 3:  # 右侧
                c = s + w0, s, s + w0 + w, s + h
            elif i == 4:  # 右下角
                c = s + w0, s + hp, s + w0 + w, s + hp + h
            elif i == 5:  # 下方
                c = s + w0 - w, s + h0, s + w0, s + h0 + h
            elif i == 6:  # 左下角
                c = s + w0 - wp - w, s + h0, s + w0 - wp, s + h0 + h
            elif i == 7:  # 左侧
                c = s - w, s + h0 - h, s, s + h0
            elif i == 8:  # 左上角
                c = s - w, s + h0 - hp - h, s, s + h0 - hp

            padw, padh = c[:2]
            x1, y1, x2, y2 = (max(x, 0) for x in c)  # 分配坐标

            # 图像
            img9[y1:y2, x1:x2] = img[y1 - padh :, x1 - padw :]  # img9[ymin:ymax, xmin:xmax]
            hp, wp = h, w  # 上一次的高度、宽度,用于下一次迭代

            # 假设 imgsz*2 的马赛克尺寸
            labels_patch = self._update_labels(labels_patch, padw + self.border[0], padh + self.border[1])
            mosaic_labels.append(labels_patch)
        final_labels = self._cat_labels(mosaic_labels)

        final_labels["img"] = img9[-self.border[0] : self.border[0], -self.border[1] : self.border[1]]
        return final_labels

    @staticmethod
    def _update_labels(labels, padw, padh):
        """Update labels.
        更新标签。
        """
        nh, nw = labels["img"].shape[:2]
        labels["instances"].convert_bbox(format="xyxy")
        labels["instances"].denormalize(nw, nh)
        labels["instances"].add_padding(padw, padh)
        return labels

    def _cat_labels(self, mosaic_labels):
        """Return labels with mosaic border instances clipped.
        返回带有马赛克边界实例剪裁的标签。
        """
        if len(mosaic_labels) == 0:
            return {}
        cls = []
        instances = []
        imgsz = self.imgsz * 2  # 马赛克图像尺寸
        for labels in mosaic_labels:
            cls.append(labels["cls"])
            instances.append(labels["instances"])
        # 最终标签
        final_labels = {
            "im_file": mosaic_labels[0]["im_file"],
            "ori_shape": mosaic_labels[0]["ori_shape"],
            "resized_shape": (imgsz, imgsz),
            "cls": np.concatenate(cls, 0),
            "instances": Instances.concatenate(instances, axis=0),
            "mosaic_border": self.border,
        }
        final_labels["instances"].clip(imgsz, imgsz)
        good = final_labels["instances"].remove_zero_area_boxes()
        final_labels["cls"] = final_labels["cls"][good]
        if "texts" in mosaic_labels[0]:
            final_labels["texts"] = mosaic_labels[0]["texts"]
        return final_labels

这个 Mosaic 类通过组合多个图像(4 或 9 个)创建一个马赛克图像,并更新对应的标签信息。

  • 初始化:我们需要传入数据集对象、图像大小、应用增强的概率和马赛克的网格大小(4 或 9)。如果概率不是在 0 到 1 之间,或者网格大小不是 4 或 9,就会报错。然后设置边界的大小为图像大小的一半。
  • 获取索引:我们根据是否使用缓冲区来决定是从缓冲区中选择图像还是从整个数据集中随机选择图像。这样可以控制从哪里选择图像来进行马赛克增强。
  • 混合变换:根据网格大小(3、4 或 9)选择相应的马赛克变换方法。如果网格大小是 3,则选择 _mosaic3 方法,如果是 4,则选择 _mosaic4 方法(其中_mosaic3 仅是给出了代码实际没有使用到,同时大部分都是采用_mosaic4 ),依此类推。
  • 马赛克变换:对于每种网格大小,我们都有对应的方法来创建马赛克图像:
    • _mosaic3:将 3 张图像按左、中、右排列。
    • _mosaic4:将 4 张图像按左上、右上、左下、右下排列。
    • _mosaic9:将 9 张图像按中心、上、右上、右、右下、下、左下、左、左上排列。
  • 更新标签:在每个马赛克变换方法中,我们都需要更新标签的信息,比如转换边界框的格式、去归一化、添加填充等。
  • 合并标签:最后,我们将所有的标签合并成一个,并处理边界问题,去除无效区域。这样我们就得到了最终的马赛克图像及其对应的标签。

这个类的主要功能是通过组合多个图像形成一个马赛克图像,从而增加数据的多样性,帮助模型更好地泛化。


三、图片示例 

上面的代码解释大家如果没有一定的基础可能看的云里雾里,下面我们经过图片示例的操作帮助大家来理解。

假设我们的网络模型首先生成一个640x640的空白图。

之后我们生成一个 (R114,G114,B114)底图,如下所示!

​ 

之后在其中选取一点最为之后四张图片的拼接点,假设其为点C坐标为(224,224) 如下图所示(这个点就是我们之后四张图片的拼接点,我们下面拿图片来举例)!

​ 

之后我们从缓存区中随机选取出四张图片或者从数据集中随机选取按照下图进行填充。

​ 

之后我们根据最开始生成白底将多余的图片内容清除掉!

 

上面的图片我们就生成了一张简单的Mosaic增强图片,其中的拼接点实际中是没有缝隙的我这里方便展示没有完全覆盖掉,下面是一张经过模型真正处理的Mosaic图片! 


四、 Mosaic 数据增强的优点和缺点

优点

1. 增加数据多样性:
   Mosaic 技术通过将多张图像拼接成一个新的图像,创造出更多样化的训练样本。这对于增强模型的泛化能力非常有帮助,特别是在数据量有限的情况下。

2. 提升模型鲁棒性:
   由于 Mosaic 图像中包含多个不同的场景和对象,模型能够在一个训练样本中学习到更多的特征,从而提高在不同场景下的鲁棒性。

3. 改进小目标检测性能:
   Mosaic 增强可以通过增加图像中小目标的数量,提升模型对小目标的检测能力。多个图像拼接后,原本单独存在的小目标在新图像中可能会变得更明显。

4. 更高的训练效率:
   通过在一个样本中包含多个场景和对象,Mosaic 增强可以在相同数量的训练步骤内让模型看到更多的数据,提升训练效率。

缺点

1. 复杂度增加:
   Mosaic 增强需要额外的处理步骤,例如图像调整大小、拼接和标签更新。这些步骤增加了数据处理的复杂性和计算开销。

2. 潜在的图像质量问题:
   由于图像被调整大小并拼接在一起,可能会导致图像质量下降,尤其是当图像被拉伸或压缩时。这可能会对模型的学习产生负面影响。

3. 标签处理复杂:
   拼接图像后,目标的位置和大小需要重新计算,这增加了标签处理的复杂性。错误的标签更新可能会导致模型训练出现问题。

4. 可能的训练不稳定性:
   在某些情况下,Mosaic 增强可能会引入不稳定性,特别是在目标分布不均的情况下。例如,如果拼接后的图像中目标过于集中或稀疏,可能会影响训练效果。

以上内容部分在网络总结,为经过我的筛选和总结出来的内容!

 本系列下一节将会更新YOLOv8损失函数的讲解!


五、本文总结

到此本文的正式分享内容就结束了,在这里给大家推荐我的YOLOv10改进有效涨点专栏,本专栏目前为新开的平均质量分98分,后期我会根据各种最新的前沿顶会进行论文复现,也会对一些老的改进机制进行补充,目前本专栏免费阅读(暂时,大家尽早关注不迷路~),如果大家觉得本文帮助到你了,订阅本专栏,关注后续更多的更新~

 专栏回顾:YOLOv10改进系列专栏——本专栏持续复习各种顶会内容——科研必备 

d2e5d4828bd84bc79d11a9bd3ef13a35.png

;