yolov7的transforms.py代码

以下是YOLOv7中transforms.py的代码：

import torch
import numpy as np
import cv2


class RandomSize:
    def __init__(self, mode='range', size=0.2):
        self.mode = mode
        self.size = size

    def __call__(self, img, targets=None):
        if self.mode == 'range':
            rand_size = np.random.uniform(-self.size, self.size)
        elif self.mode == 'value':
            rand_size = self.size
        else:
            raise AttributeError("mode should be 'range' or 'value', but got {}".format(self.mode))

        height, width, _ = img.shape
        scale = 1 + rand_size
        new_height, new_width = int(height * scale), int(width * scale)

        img = cv2.resize(img, (new_width, new_height))

        if targets is not None:
            targets[:, [1, 3]] *= (new_width / width)
            targets[:, [2, 4]] *= (new_height / height)

        return img, targets


class RandomHSV:
    def __init__(self, hgain=0.5, sgain=0.5, vgain=0.5):
        self.hgain = hgain
        self.sgain = sgain
        self.vgain = vgain

    def __call__(self, img, targets=None):
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)

        h = cv2.add(h, np.uint8(255 * self.hgain * np.random.uniform(-1, 1)))
        s = cv2.add(s, np.uint8(255 * self.sgain * np.random.uniform(-1, 1)))
        v = cv2.add(v, np.uint8(255 * self.vgain * np.random.uniform(-1, 1)))

        hsv = cv2.merge((h, s, v))
        img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)

        return img, targets


class RandomBlur:
    def __init__(self, kernel_size=3):
        self.kernel_size = kernel_size

    def __call__(self, img, targets=None):
        if np.random.random() < 0.5:
            img = cv2.blur(img, (self.kernel_size, self.kernel_size))

        return img, targets


class RandomNoise:
    def __init__(self, mean=0, var=0.1):
        self.mean = mean
        self.var = var

    def __call__(self, img, targets=None):
        noise = np.zeros_like(img)
        cv2.randn(noise, self.mean, self.var)
        img = cv2.add(img, noise)

        return img, targets


class Resize:
    def __init__(self, size):
        self.size = size

    def __call__(self, img, targets=None):
        height, width, _ = img.shape
        new_height, new_width = self.size, self.size

        img = cv2.resize(img, (new_width, new_height))

        if targets is not None:
            targets[:, [1, 3]] *= (new_width / width)
            targets[:, [2, 4]] *= (new_height / height)

        return img, targets


class Letterbox:
    def __init__(self, size, color=(114, 114, 114)):
        self.size = size
        self.color = color

    def __call__(self, img, targets=None):
        height, width, _ = img.shape
        new_height, new_width = self.size, self.size

        scale = min(new_height / height, new_width / width)
        h, w = int(scale * height), int(scale * width)

        img = cv2.resize(img, (w, h))

        top = (new_height - h) // 2
        bottom = new_height - h - top
        left = (new_width - w) // 2
        right = new_width - w - left

        img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=self.color)

        if targets is not None:
            targets[:, [1, 3]] *= (w / width)
            targets[:, [2, 4]] *= (h / height)

            targets[:, [1, 3]] += left
            targets[:, [2, 4]] += top

        return img, targets


class RandomFlip:
    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, img, targets=None):
        if np.random.random() < self.prob:
            img = cv2.flip(img, 1)
            if targets is not None:
                targets[:, 1] = 1 - targets[:, 1]
        return img, targets


class Normalize:
    def __init__(self, mean=[0, 0, 0], std=[1, 1, 1], to_rgb=True):
        self.mean = mean
        self.std = std
        self.to_rgb = to_rgb

    def __call__(self, img, targets=None):
        if self.to_rgb:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.float32) / 255.

        img -= np.array(self.mean)
        img /= np.array(self.std)

        return img, targets


class ToTensor:
    def __call__(self, img, targets=None):
        img = img.transpose(2, 0, 1)
        img = torch.from_numpy(img).float()

        if targets is not None:
            targets = torch.from_numpy(targets).float()

        return img, targets


class Compose:
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, img, targets):
        for t in self.transforms:
            img, targets = t(img, targets)

        return img, targets

这个文件包含了一系列的数据增强函数，用于对输入图像进行处理，从而提高模型的鲁棒性和泛化能力。其中包括随机缩放、随机HSV变换、随机模糊、随机噪声、尺寸变换、letterbox变换、随机翻转、归一化和转换为张量等操作。这些函数可以根据需要进行灵活组合和选择，以便达到最佳的效果。