본문 바로가기

Deep Learning

Image Augmentation (Photometric) 방법

반응형
class PhotoMetricDistortion:
    """Apply photometric distortion to image sequentially, every transformation
    is applied with a probability of 0.5. The position of random contrast is in
    second or second to last.
    1. random brightness
    2. random contrast (mode 0)
    3. convert color from BGR to HSV
    4. random saturation
    5. random hue
    6. convert color from HSV to BGR
    7. random contrast (mode 1)
    8. randomly swap channels
    Args:
        brightness_delta (int): delta of brightness.
        contrast_range (tuple): range of contrast.
        saturation_range (tuple): range of saturation.
        hue_delta (int): delta of hue.
    """

    def __init__(self,
                 brightness_delta=32,
                 contrast_range=(0.5, 1.5),
                 saturation_range=(0.5, 1.5),
                 hue_delta=9):
        self.brightness_delta = brightness_delta
        self.contrast_lower, self.contrast_upper = contrast_range
        self.saturation_lower, self.saturation_upper = saturation_range
        self.hue_delta = hue_delta
        self.dtype = np.float32

    def __call__(self, img):
        """
        img : image of type PIL.Image
        """

        # conversion to np.array of type np.float32
        img = np.array(img).astype(self.dtype)


        # random brightness
        if random.randint(2):
            delta = random.uniform(-self.brightness_delta, self.brightness_delta)
            img = cv2.cvtColor(img.astype('uint8'), cv2.COLOR_RGB2YCrCb).astype(self.dtype)
            img[..., 0] += delta
            img[img[..., 0] > 255, 0] = 255.0
            img = cv2.cvtColor(img.astype('uint8'), cv2.COLOR_YCrCb2RGB).astype(self.dtype)
            # print(f'random brightness: {delta}')


        # mode == 0 --> do random contrast first
        # mode == 1 --> do random contrast last
        mode = random.randint(2)
        if mode == 1:
            if random.randint(2):
                alpha = random.uniform(self.contrast_lower, self.contrast_upper)
                img *= alpha
                img[img > 255.0] = 255.0
                # print(f'random contrast: {alpha}')


        # convert color from BGR to HSV
        img = cv2.cvtColor(img.astype('uint8'), cv2.COLOR_RGB2HSV).astype(self.dtype)


        # random saturation
        if random.randint(2):
            alpha = random.uniform(self.saturation_lower, self.saturation_upper)
            img[..., 1] *= alpha
            img[img[..., 1] > 255.0] = 255.0
            # print(f'random saturation: {alpha}')


        # random hue
        if random.randint(2):
            delta = random.uniform(-self.hue_delta, self.hue_delta)
            img[..., 0] += delta
            img[..., 0][img[..., 0] > 179] = 179
            # print(f'random hue: {delta}')

        # convert color from HSV to BGR
        img = cv2.cvtColor(img.astype('uint8'), cv2.COLOR_HSV2RGB).astype(self.dtype)

        # random contrast
        if mode == 0:
            if random.randint(2):
                alpha = random.uniform(self.contrast_lower, self.contrast_upper)
                img *= alpha
                # print(f'random contrast: {alpha}')

        # randomly swap channels
        if random.randint(2):
            img = img[..., random.permutation(3)]
            # print('random permutation')

        return Image.fromarray(img.astype('uint8'))