YOLOv5数据增强算法替代

Question

YOLOv5-7.0版本的算法中，数据增强方式采取了Mosaic方法，并且在Mosaic被采用的时候才可以调用mixup数据增强方法。那怎样可以用FMix算法替换mixup，以便使YOLOv5算法的数据增强效果更强大？

Answer 1

“Devil组”引证GPT后的撰写：

import torch
import numpy as np
import albumentations as A

def fmix(x, y, alpha=1.0, decay_power=3.0, shape=(416, 416), max_soft=0.3, reformulate=False):
    lam, mask = sample_mask(alpha, decay_power, shape, max_soft)
    index = torch.randperm(x.size()[0]).cuda()
    xf = lam * x + (1 - lam) * x[index, :]
    yf = lam * y + (1 - lam) * y[index, :]
    if reformulate:
        yf[:, 0:4] = bbox_mashup(yf[:, 0:4], y[index, :, 0:4], mask)

    return xf, yf, lam

def sample_mask(alpha, decay_power, shape, max_soft):
    H, W = shape
    lam, mask = np.random.beta(alpha, alpha), np.zeros((H, W))
    indices = np.random.choice(np.arange(H*W), int(H*W*lam), replace=False)
    mask[np.unravel_index(indices, (H, W))] = 1
    frequency = cv2.GaussianBlur(mask, (21, 21), 0)
    frequency = (frequency - frequency.min()) / (frequency.max() - frequency.min())
    frequency = (max_soft - 0.1) * frequency + 0.1
    return lam, frequency

def bbox_mashup(src_bbox, dst_bbox, mask):
    src_bbox[:, :2] = mask * src_bbox[:, :2] + (1 - mask) * dst_bbox[:, :2]
    src_bbox[:, 2:4] = mask * src_bbox[:, 2:4] + (1 - mask) * dst_bbox[:, 2:4]
    return src_bbox
class YOLOv5Dataset(torch.utils.data.Dataset):
    def __init__(self, data, img_size=416, transform=None, mosaic=False, mixup=False, fmix=False):
        self.img_files = []
        self.label_files = []
        self.img_size = img_size
        self.transform = transform
        self.mosaic = mosaic
        self.mixup = mixup
        self.fmix = fmix

        for d in data:
            if isinstance(d, str):
                if os.path.isdir(d):
                    self.img_files += glob.glob(os.path.join(d, '*.jpg'))
                else:
                    self.img_files.append(d)
            else:
                self.img_files.append(d[0])
                self.label_files.append(d[1])

    def __getitem__(self, index):
        if self.mixup:
            img, label = self._mixup(index)
        elif self.mosaic:
            img, label = self._mosaic(index)
        elif self.fmix:
            img, label = self._fmix(index)
        else:
            img, label = self._get_item(index)

        if self.transform:
            img, label = self.transform(img, label)

        return img, label
def _get_item(self, index):
    img_path = self.img_files[index]
    img = cv2.imread(img_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    h, w, _ = img.shape
    label_path = self.label_files[index]
    label = []

    if os.path.exists(label_path):
        with open(label_path, 'r') as f:
            lines = f.readlines()
            for line in lines:
                line = line.strip()
                if len(line) > 0:
                    class_id, x, y, w, h = line.split()
                    x, y, w, h = float(x), float(y), float(w), float(h)
                    label.append([x, y, w, h, int(class_id)])

    if len(label) == 0:
        label.append([-1, -1, -1, -1, -1])

    label = np.array(label)

    return img, label
if fmix:
    x, y, lam = fmix(x, y)