python实现 YOLO中的数据增强|Mixup、Mosaic

Mixup

Mixup指的是随机在训练集中抽取两类，用线性插值方法将他们混合在一起。
假设当前文件夹下共有四张图片

import os
import random
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

# 文件夹路径和图像文件名
image_folder = './'
image_files = ['image1.jpg', 'image2.jpg', 'image3.jpg', 'image4.jpg']

# 加载图片
def load_image(path):
    img = Image.open(path).convert("RGB")
    img = img.resize((640, 640))
    img = np.array(img) / 255.0
    return img


images = [load_image(os.path.join(image_folder, img)) for img in image_files]

# 随机生成0~1的lam，再融合
def mixup_data(images,alpha=1.0):
    # Beta 分布通常用于生成介于 0 和 1 之间的随机数。
    lam = np.random.beta(alpha,alpha)
    idx = random.sample(range(len(images)),2)
    image1 ,image2 = images[idx[0]],images[idx[1]]
    mixed = lam*image1 + (1-lam)*image2
    return mixed,lam

mixed_images,lam = mixup_data(images)
def show_images(image):
    plt.imshow(image)
    plt.axis('off')
    plt.show()
    
show_images(mixed_images)    

运行后：

Mosaic

比起copy paste方法要将实例从图片中抠出来，在YoloV4方法中提出的Mosaic数据增强更加简单，先将四张图片做各种变换（比如翻转、调整尺寸等），再将他们拼贴在一起，作为网络的输入。这种方法不仅能提升模型的鲁棒性，也能减少mini-batchsize，加快训练的速度。

import os
import random
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

# 文件夹路径和图像文件名
image_folder = './'
image_files = ['image1.jpg', 'image2.jpg', 'image3.jpg', 'image4.jpg']

def load_image(path):
    img = Image.open(path).convert("RGB")
    img = img.resize((640, 640))
    img = np.array(img) / 255.0
    return img

# 加载图像
images = [load_image(os.path.join(image_folder, img)) for img in image_files]

"""对图像进行随机变换"""
def random_transform(image):
    img = Image.fromarray((image * 255).astype(np.uint8))
    # 随机水平翻转
    if random.random() > 0.5:
        img = img.transpose(Image.FLIP_LEFT_RIGHT)
    
    # 随机垂直翻转
    if random.random() > 0.5:
        img = img.transpose(Image.FLIP_TOP_BOTTOM)
    
    # 随机旋转
    angle = random.choice([0, 90, 180, 270])
    img = img.rotate(angle)
    
    img = img.resize((640, 640))
    return np.array(img) / 255.0

def random_crop(image, crop_size):
    """从图像中随机裁剪一个区域"""
    h, w, _ = image.shape
    top = np.random.randint(0, h - crop_size)
    left = np.random.randint(0, w - crop_size)
    cropped_image = image[top:top + crop_size, left:left + crop_size]
    return cropped_image
    
def mosaic(images, output_size=(640, 640)):
    
    assert len(images) == 4
    
    # 定义输出图像尺寸
    output_image = np.zeros((output_size[0], output_size[1], 3))
    
    # 随机生成切割点
    cut_x = random.randint(int(output_size[0] * 0.3), int(output_size[0] * 0.7))
    cut_y = random.randint(int(output_size[1] * 0.3), int(output_size[1] * 0.7))
    
    # 变换四张图像
    transformed_images = [random_transform(img) for img in images]
    
    # 将四个变换后的图像分别放置在新图像的四个象限
    output_image[:cut_y, :cut_x] = transformed_images[0][:cut_y, :cut_x]
    output_image[:cut_y, cut_x:] = transformed_images[1][:cut_y, cut_x:]
    output_image[cut_y:, :cut_x] = transformed_images[2][cut_y:, :cut_x]
    output_image[cut_y:, cut_x:] = transformed_images[3][cut_y:, cut_x:]
    
    return output_image
    
# 进行 Mosaic 增强
augmented_image = mosaic(images)

# 显示增强后的图像
plt.imshow(augmented_image)
plt.axis('off')
plt.show()

运行之后：