ImageOnlyTransformを継承したRandomShadowの実装

random_shadow.py

import math
import random
from albumentations.core.transforms_interface import ImageOnlyTransform


class RandomShadow(ImageOnlyTransform):
    def __init__(self, always_apply=False, p=0.5, sl=0.02, sh=0.4, r1=0.3, brightness=0.75):
        super(RandomShadow, self).__init__(always_apply=always_apply, p=p)
        self.brightness = brightness
        self.sl = sl
        self.sh = sh
        self.r1 = r1
    def apply(self, img, **params):
        for attempt in range(100):
            area = img.shape[0] * img.shape[1]
            target_area = random.uniform(self.sl, self.sh) * area
            aspect_ratio = random.uniform(self.r1, 1 / self.r1)
            h = int(round(math.sqrt(target_area * aspect_ratio)))
            w = int(round(math.sqrt(target_area / aspect_ratio)))
            if w < img.shape[1] and h < img.shape[0]:
                x1 = random.randint(0, img.shape[0] - h)
                y1 = random.randint(0, img.shape[1] - w)
                if len(img.shape) == 3:
                    img[x1:x1 + h, y1:y1 + w, 0] = img[x1:x1 + h, y1:y1 + w, 0] * self.brightness
                    img[x1:x1 + h, y1:y1 + w, 1] = img[x1:x1 + h, y1:y1 + w, 1] * self.brightness
                    img[x1:x1 + h, y1:y1 + w, 2] = img[x1:x1 + h, y1:y1 + w, 2] * self.brightness
                else:
                    img[x1:x1 + h, y1:y1 + w] = img[x1:x1 + h, y1:y1 + w] * self.brightness
                return img
        return img