yolov5实现数据增强

Yolov5是一种目标检测算法，数据增强是提高模型鲁棒性和泛化能力的重要手段。以下是Yolov5中常用的数据增强方式及相应的源代码：

1. 随机缩放

from PIL import Image
import random

def random_resize(img, targets, size_range=(0.2, 2)):
    """
    随机缩放图像和目标框
    :param img: PIL.Image对象
    :param targets: 目标框列表，每个目标框为[x1, y1, x2, y2, class]
    :param size_range: 缩放比例范围
    :return: 缩放后的图像和目标框
    """
    w, h = img.size
    size_ratio = random.uniform(*size_range)
    new_w, new_h = int(w * size_ratio), int(h * size_ratio)
    img = img.resize((new_w, new_h), resample=Image.BILINEAR)
    new_targets = []
    for target in targets:
        x1, y1, x2, y2, cls = target
        x1, x2 = x1 * size_ratio, x2 * size_ratio
        y1, y2 = y1 * size_ratio, y2 * size_ratio
        new_targets.append([x1, y1, x2, y2, cls])
    return img, new_targets

2. 随机裁剪

import numpy as np

def random_crop(img, targets, crop_size=(640, 640)):
    """
    随机裁剪图像和目标框
    :param img: PIL.Image对象
    :param targets: 目标框列表，每个目标框为[x1, y1, x2, y2, class]
    :param crop_size: 裁剪尺寸
    :return: 裁剪后的图像和目标框
    """
    w, h = img.size
    crop_w, crop_h = crop_size
    if w == crop_w and h == crop_h:
        return img, targets
    if w < crop_w or h < crop_h:
        img = img.resize((crop_w, crop_h), resample=Image.BILINEAR)
        new_targets = []
        for target in targets:
            x1, y1, x2, y2, cls = target
            x1, x2 = x1 * crop_w / w, x2 * crop_w / w
            y1, y2 = y1 * crop_h / h, y2 * crop_h / h
            new_targets.append([x1, y1, x2, y2, cls])
        return img, new_targets
    x = np.random.randint(0, w - crop_w)
    y = np.random.randint(0, h - crop_h)
    img = img.crop((x, y, x + crop_w, y + crop_h))
    new_targets = []
    for target in targets:
        x1, y1, x2, y2, cls = target
        x1, x2 = x1 - x, x2 - x
        y1, y2 = y1 - y, y2 - y
        if x1 < crop_w and x2 > 0 and y1 < crop_h and y2 > 0:
            x1 = max(0, x1)
            y1 = max(0, y1)
            x2 = min(crop_w, x2)
            y2 = min(crop_h, y2)
            new_targets.append([x1, y1, x2, y2, cls])
    return img, new_targets

3. 随机翻转

def random_flip(img, targets):
    """
    随机翻转图像和目标框
    :param img: PIL.Image对象
    :param targets: 目标框列表，每个目标框为[x1, y1, x2, y2, class]
    :return: 翻转后的图像和目标框
    """
    if random.random() < 0.5:
        img = img.transpose(Image.FLIP_LEFT_RIGHT)
        new_targets = []
        for target in targets:
            x1, y1, x2, y2, cls = target
            new_targets.append([img.width - x2, y1, img.width - x1, y2, cls])
        return img, new_targets
    return img, targets