4. Yolov8 Mosaic数据增强详解

Mosaic数据增强简介

        Mosaic数据增强是YOLOV4论文中提出的一种数据增强方法。其主要思想是将四张图片进行随机裁剪，然后拼接到一张图上作为训练数据。这样做的好处是丰富了图片的背景，并且四张图片拼接在一起变相地提高了batch_size，在进行batch normalization的时候也会计算四张图片，所以对本身batch_size不是很依赖。

Mosaic数据增强方法的实现过程大致分为三步：

        1.从数据集中每次随机读取四张图片。
        2.分别对四张图片进行翻转（对原始图片进行左右的翻转）、缩放（对原始图片进行大小的缩放）、色域变化（对原始图片的明亮度、饱和度、色调进行改变）等操作。操作完成之后然后再将原始图片按照 第一张图片摆放在左上，第二张图片摆放在左下，第三张图片摆放在右下，第四张图片摆放在右上四个方向位置摆好。
        3.进行图片的组合和框的组合。完成四张图片的摆放之后，我们利用矩阵的方式将四张图片它固定的区域截取下来，然后将它们拼接起来，拼接成一 张新的图片，新的图片上含有框框等一系列的内容。
        Mosaic数据增强方法参考了CutMix数据增强方式, 是CutMix数据增强方法的改进版。CutMix算法使用两张图片进行拼接，而Mosaic算法一般使用四张进行拼接。这种方法极大地丰富了检测物体的背景，并且在标准化BN计算的时候一下子会计算四张图片的数据。这种方法可以有效地提高模型的性能和鲁棒性。

代码及详解

import random
import numpy as np

class BaseMixTransform:
    """This implementation is from mmyolo"""

    def __init__(self, dataset, pre_transform=None, p=0.0) -> None:
        # dataset中存储着读取的数据（图片、类别和真实框位置信息）
        self.dataset = dataset
        # 进行数据增强前进行的数据预处理
        self.pre_transform = pre_transform
        # 进行数据增强的概率
        self.p = p

    def __call__(self, labels):
        # 随机进行数据增强
        if random.uniform(0, 1) > self.p:
            return labels

        # 获取一个或三个其他图像的索引
        indexes = self.get_indexes()
        if isinstance(indexes, int):
            indexes = [indexes]

        # 获取图像信息用于 Mosaic 或者 MixUp 数据增强
        mix_labels = [self.dataset.get_label_info(i) for i in indexes]

        if self.pre_transform is not None:
            for i, data in enumerate(mix_labels):
                mix_labels[i] = self.pre_transform(data)
        labels["mix_labels"] = mix_labels

        # 进行 Mosaic 或 MixUp
        labels = self._mix_transform(labels)
        labels.pop("mix_labels", None)
        return labels

    def _mix_transform(self, labels):
        raise NotImplementedError

    def get_indexes(self):
        raise NotImplementedError


class Mosaic(BaseMixTransform):
    """Mosaic augmentation.
    Args:
        imgsz (Sequence[int]): Image size after mosaic pipeline of single
            image. The shape order should be (height, width).
            Default to (640, 640).
    """

    def __init__(self, dataset, imgsz=640, p=1.0, border=(0, 0)):
        assert 0 <= p <= 1.0, "The probability should be in range [0, 1]. " f"got {p}."
        super().__init__(dataset=dataset, p=p)
        self.dataset = dataset
        self.imgsz = imgsz
        self.border = border
    
    # 获取用于数据增强的其它三个图片的索引
    def get_indexes(self):
        return [random.randint(0, len(self.dataset) - 1) for _ in range(3)]

    def _mix_transform(self, labels):
        mosaic_labels = []
        assert labels.get("rect_shape", None) is None, "rect and mosaic is exclusive."
        assert len(labels.get("mix_labels", [])) > 0, "There are no other images for mosaic augment."
        s = self.imgsz
        # random.uniform 生成一个在 -x ~ 2*s+x 的浮点数
        # 随机选出 xc 和 yc，（xc，yc）是大图的中心位置。
        yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.border)  # mosaic center x, y
        for i in range(4):
            # 获取图片所有信息（图片数据、类别和真实框位置信息）
            labels_patch = (labels if i == 0 else labels["mix_labels"][i - 1]).copy()
            # 从图片的所有信息中取出图片数据
            img = labels_patch["img"]
            # 从图片的所有信息中取出预处理后图片的位置信息
            h, w = labels_patch.pop("resized_shape")

            # 把四张图片分别放到 img4(Mosaic数据增强后的大图)
            if i == 0:  # top left
                # 生成一个(2*imgsz, 2*imgsz, 3)的大图，将四张图片依次铺到此图上。
                img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles
                # 确定小图将要放到大图左上的位置
                x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)
                # 获取小图的大小
                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
                x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
            elif i == 2:  # bottom left
                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:  # bottom right
                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 _update_labels(self, labels, padw, padh):
        """Update labels"""
        nh, nw = labels["img"].shape[:2]
        # xywh -> xyxy
        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):
        if len(mosaic_labels) == 0:
            return {}
        cls = []
        instances = []
        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": (self.imgsz * 2, self.imgsz * 2),
            "cls": np.concatenate(cls, 0),
            "instances": Instances.concatenate(instances, axis=0),
            "mosaic_border": self.border}
        final_labels["instances"].clip(self.imgsz * 2, self.imgsz * 2)
        return final_labels