unilm/trocr/augmentation/weather.py

import cv2
import numpy as np
import math
from PIL import Image, ImageOps, ImageDraw
from skimage import color
from scipy import interpolate
from pkg_resources import resource_filename
from io import BytesIO
from .ops import plasma_fractal, clipped_zoom, MotionImage

'''
    PIL resize (W,H)
'''
class Fog:
    def __init__(self):
        pass

    def __call__(self, img, mag=-1, prob=1.):
        if np.random.uniform(0,1) > prob:
            return img

        W, H = img.size
        c = [(1.5, 2), (2., 2), (2.5, 1.7)]
        if mag<0 or mag>=len(c):
            index = np.random.randint(0, len(c))
        else:
            index = mag
        c = c[index]

        n_channels = len(img.getbands())
        isgray = n_channels == 1

        img = np.array(img) / 255.
        max_val = img.max()
        fog = c[0] * plasma_fractal(wibbledecay=c[1])[:H, :W][..., np.newaxis]
        #x += c[0] * plasma_fractal(wibbledecay=c[1])[:224, :224][..., np.newaxis]
        #return np.clip(x * max_val / (max_val + c[0]), 0, 1) * 255
        if isgray:
            fog = np.squeeze(fog)
        else:
            fog = np.repeat(fog, 3, axis=2)

        # print('img', img.shape)
        # print('fog', fog.shape)
        # print(H, W)
        # exit(0)
        fog = cv2.resize(fog, dsize=(H, W), interpolation=cv2.INTER_CUBIC)
        img += fog
        img = np.clip(img * max_val / (max_val + c[0]), 0, 1) * 255
        return Image.fromarray(img.astype(np.uint8))


class Frost:
    def __init__(self):
        pass

    def __call__(self, img, mag=-1, prob=1.):
        if np.random.uniform(0,1) > prob:
            return img

        W, H = img.size
        c = [(1, 0.4), (0.8, 0.6), (0.7, 0.7)]
        if mag<0 or mag>=len(c):
            index = np.random.randint(0, len(c))
        else:
            index = mag
        c = c[index]

        filename = [resource_filename(__name__, 'frost/frost1.png'),
                    resource_filename(__name__, 'frost/frost2.png'),
                    resource_filename(__name__, 'frost/frost3.png'),
                    resource_filename(__name__, 'frost/frost4.jpg'),
                    resource_filename(__name__, 'frost/frost5.jpg'),
                    resource_filename(__name__, 'frost/frost6.jpg')
                    ]
        index = np.random.randint(0, len(filename))
        filename = filename[index]
        frost = cv2.imread(filename)
        frost = cv2.resize(frost, dsize=(H, W), interpolation=cv2.INTER_CUBIC)
        #randomly crop and convert to rgb
        # x_start, y_start = np.random.randint(0, frost.shape[0] - H), np.random.randint(0, frost.shape[1] - W)
        x_start = 0
        y_start = 0
        frost = frost[x_start:x_start + H, y_start:y_start + W][..., [2, 1, 0]]

        n_channels = len(img.getbands())
        isgray = n_channels == 1

        img = np.array(img)
        
        if isgray:
            img = np.expand_dims(img, axis=2)
            img = np.repeat(img, 3, axis=2)

        img = img * c[0]
        frost = frost * c[1]
        img = np.clip(c[0] * img + c[1] * frost, 0, 255)
        img = Image.fromarray(img.astype(np.uint8))
        if isgray:
            img = ImageOps.grayscale(img)

        return img

class Snow:
    def __init__(self):
        pass

    def __call__(self, img, mag=-1, prob=1.):
        if np.random.uniform(0,1) > prob:
            return img

        W, H = img.size
        c = [(0.1, 0.3, 3, 0.5, 10, 4, 0.8),
             (0.2, 0.3, 2, 0.5, 12, 4, 0.7),
             (0.55, 0.3, 4, 0.9, 12, 8, 0.7)]
        if mag<0 or mag>=len(c):
            index = np.random.randint(0, len(c))
        else:
            index = mag
        c = c[index]

        n_channels = len(img.getbands())
        isgray = n_channels == 1

        img = np.array(img, dtype=np.float32) / 255.
        if isgray:
            img = np.expand_dims(img, axis=2)
            img = np.repeat(img, 3, axis=2)

        snow_layer = np.random.normal(size=img.shape[:2], loc=c[0], scale=c[1])  # [:2] for monochrome

        #snow_layer = clipped_zoom(snow_layer[..., np.newaxis], c[2])
        snow_layer[snow_layer < c[3]] = 0

        snow_layer = Image.fromarray((np.clip(snow_layer.squeeze(), 0, 1) * 255).astype(np.uint8), mode='L')
        output = BytesIO()
        snow_layer.save(output, format='PNG')
        snow_layer = MotionImage(blob=output.getvalue())

        snow_layer.motion_blur(radius=c[4], sigma=c[5], angle=np.random.uniform(-135, -45))

        snow_layer = cv2.imdecode(np.fromstring(snow_layer.make_blob(), np.uint8),
                                  cv2.IMREAD_UNCHANGED) / 255.

        #snow_layer = cv2.cvtColor(snow_layer, cv2.COLOR_BGR2RGB)

        snow_layer = snow_layer[..., np.newaxis]

        img = c[6] * img
        gray_img = (1 - c[6]) * np.maximum(img, cv2.cvtColor(img, cv2.COLOR_RGB2GRAY).reshape(H, W, 1) * 1.5 + 0.5)
        img += gray_img
        img = np.clip(img + snow_layer + np.rot90(snow_layer, k=2), 0, 1) * 255
        img = Image.fromarray(img.astype(np.uint8))
        if isgray:
            img = ImageOps.grayscale(img)

        return img

class Rain:
    def __init__(self):
        pass

    def __call__(self, img, mag=-1, prob=1.):
        if np.random.uniform(0,1) > prob:
            return img

        img = img.copy()
        W, H = img.size
        n_channels = len(img.getbands())
        isgray = n_channels == 1
        line_width = np.random.randint(1, 2)

        c =[50, 70, 90]
        if mag<0 or mag>=len(c):
            index = 0
        else:
            index = mag
        c = c[index]

        n_rains = np.random.randint(c, c+20)
        slant = np.random.randint(-60, 60)
        fillcolor = 200 if isgray else (200,200,200)

        draw = ImageDraw.Draw(img)
        for i in range(1, n_rains):
            length = np.random.randint(5, 10)
            x1 = np.random.randint(0, W-length)
            y1 = np.random.randint(0, H-length)
            x2 = x1 + length*math.sin(slant*math.pi/180.)
            y2 = y1 + length*math.cos(slant*math.pi/180.)
            x2 = int(x2)
            y2 = int(y2)
            draw.line([(x1,y1), (x2,y2)], width=line_width, fill=fillcolor)

        return img

class Shadow:
    def __init__(self):
        pass

    def __call__(self, img, mag=-1, prob=1.):
        if np.random.uniform(0,1) > prob:
            return img

        #img = img.copy()
        W, H = img.size
        n_channels = len(img.getbands())
        isgray = n_channels == 1

        c =[64, 96, 128]
        if mag<0 or mag>=len(c):
            index = 0
        else:
            index = mag
        c = c[index]

        img = img.convert('RGBA')
        overlay = Image.new('RGBA', img.size, (255,255,255,0))
        draw = ImageDraw.Draw(overlay) 
        transparency = np.random.randint(c, c+32)
        x1 = np.random.randint(0, W//2)
        y1 = 0

        x2 = np.random.randint(W//2, W)
        y2 = 0

        x3 = np.random.randint(W//2, W)
        y3 = H - 1

        x4 = np.random.randint(0, W//2)
        y4 = H - 1

        draw.polygon([(x1,y1), (x2,y2), (x3,y3), (x4,y4)], fill=(0,0,0,transparency))

        img = Image.alpha_composite(img, overlay)
        img = img.convert("RGB")
        if isgray:
            img = ImageOps.grayscale(img)

        return img