赞
踩
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
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。