training/mono/utils/raindropper/raindrop.py

# change rainy drop func from
# https://github.com/EvoCargo/RaindropsOnWindshield/blob/main/raindrops_generator/raindrop/raindrop.py

import math
import random
from random import randint

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFilter
from raindropper.config import cfg


def make_bezier(xys):
    # xys should be a sequence of 2-tuples (Bezier control points)
    n = len(xys)
    combinations = pascal_row(n - 1)

    def bezier(ts):
        # This uses the generalized formula for bezier curves
        # http://en.wikipedia.org/wiki/B%C3%A9zier_curve#Generalization
        result = []
        for t in ts:
            tpowers = (t**i for i in range(n))
            upowers = reversed([(1 - t)**i for i in range(n)])
            coefs = [c * a * b for c, a, b in zip(combinations, tpowers, upowers)]
            result.append(tuple(sum([coef * p for coef, p in zip(coefs, ps)]) for ps in zip(*xys)))
        return result

    return bezier


def pascal_row(n, memo={}):
    # This returns the nth row of Pascal Triangle
    if n in memo:
        return memo[n]
    result = [1]
    x, numerator = 1, n
    for denominator in range(1, n // 2 + 1):
        x *= numerator
        x /= denominator
        result.append(x)
        numerator -= 1
        if n & 1 == 0:
            result.extend(reversed(result[:-1]))
        else:
            result.extend(reversed(result))
        memo[n] = result
    return result


class Raindrop():

    def __init__(self, key, centerxy=None, radius=None, shape=None):
        # param key: a unique key identifying a drop
        # param centerxy: tuple defining coordinates of raindrop center in the image
        # param radius: radius of a drop
        # param shape: int from 0 to 2 defining raindrop shape type
        self.key = key
        self.ifcol = False
        self.col_with = []
        self.center = centerxy
        self.radius = radius
        # self.blur_coeff = max(int(self.radius/3), 1)
        # self.blur_coeff = max(int(cfg["maxR"] / self.radius), 1)
        self.blur_coeff = 3
        self.shape = shape
        self.type = 'default'
        # label map's WxH = 4*R , 5*R
        self.labelmap = np.zeros((self.radius * 5, self.radius * 4))
        self.alphamap = np.zeros((self.radius * 5, self.radius * 4))
        self.background = None
        self.texture = None
        self.control_point = {}
        self._create_label()
        self.use_label = False

    def setCollision(self, col, col_with):
        self.ifcol = col
        self.col_with = col_with

    def updateTexture(self, bg):
        # gaussian blur radius may be 1, 3, 5
        radius_array = [1, 3]
        blur_radius_idx = randint(0, 1)
        blur_radius = radius_array[blur_radius_idx]
        fg = (Image.fromarray(np.uint8(bg))).filter(ImageFilter.GaussianBlur(radius=blur_radius))
        fg = np.asarray(fg)

        # add fish eye effect to simulate the background
        K = np.array([[30 * self.radius, 0, 2 * self.radius], [0., 20 * self.radius, 3 * self.radius], [0., 0., 1]])
        D = np.array([0.0, 0.0, 0.0, 0.0])
        Knew = K.copy()

        Knew[(0, 1), (0, 1)] = math.pow(self.radius, 1 / 500) * 2 * Knew[(0, 1), (0, 1)]
        fisheye = cv2.fisheye.undistortImage(fg, K, D=D, Knew=Knew)

        tmp = np.expand_dims(self.alphamap, axis=-1)
        tmp = np.concatenate((fisheye, tmp), axis=2)

        self.texture = Image.fromarray(tmp.astype('uint8'), 'RGBA')

    def _create_label(self):
        self._createDefaultDrop()

    def _createDefaultDrop(self):
        """create the raindrop Alpha Map according to its shape type update
        raindrop label."""
        if (self.shape == 0):
            cv2.circle(self.labelmap, (self.radius * 2, self.radius * 3), int(self.radius), 128, -1)
            self.alphamap = (Image.fromarray(np.uint8(self.labelmap))).filter(
                ImageFilter.GaussianBlur(radius=self.blur_coeff))
            self.alphamap = np.asarray(self.alphamap).astype(np.float)
            self.alphamap = self.alphamap / np.max(self.alphamap) * 255.0
            # set label map
            self.labelmap[self.labelmap > 0] = 1

        if (self.shape == 1):
            cv2.circle(self.labelmap, (self.radius * 2, self.radius * 3), int(self.radius), 128, -1)
            cv2.ellipse(self.labelmap, (self.radius * 2, self.radius * 3),
                        (self.radius, int(1.3 * math.sqrt(3) * self.radius)), 0, 180, 360, 128, -1)

            self.alphamap = (Image.fromarray(np.uint8(self.labelmap))).filter(
                ImageFilter.GaussianBlur(radius=self.blur_coeff))
            self.alphamap = np.asarray(self.alphamap).astype(np.float)
            self.alphamap = self.alphamap / np.max(self.alphamap) * 255.0
            # set label map
            self.labelmap[self.labelmap > 0] = 1

        if (self.shape == 2):
            C0 = random.uniform(0, 1)
            C1 = random.uniform(0, 1)
            A0 = random.uniform(0, 1)
            A1 = random.uniform(2, 3)
            D0 = random.uniform(2, 3)
            D1 = random.uniform(2, 3)
            B0 = random.uniform(2, 3)
            B1 = random.uniform(0, 1)

            self.control_point['A'] = (A0, A1)
            self.control_point['B'] = (B0, B1)
            self.control_point['C'] = (C0, C1)
            self.control_point['D'] = (D0, D1)

            img = Image.fromarray(np.uint8(self.labelmap), 'L')
            draw = ImageDraw.Draw(img)
            ts = [t / 100.0 for t in range(101)]
            xys = [(self.radius * C0, self.radius * C1), (self.radius * B0, self.radius * B1),
                   (self.radius * D0, self.radius * D1)]
            bezier = make_bezier(xys)
            points = bezier(ts)

            xys = [(self.radius * C0, self.radius * C1), (self.radius * A0, self.radius * A1),
                   (self.radius * D0, self.radius * D1)]
            bezier = make_bezier(xys)
            points.extend(bezier(ts))
            draw.polygon(points, fill='gray')

            self.alphamap = img.filter(ImageFilter.GaussianBlur(radius=self.blur_coeff))
            self.alphamap = np.asarray(self.alphamap).astype(np.float)
            self.alphamap = self.alphamap / np.max(self.alphamap) * 255.0

            # set label map
            self.labelmap[self.labelmap > 0] = 1

    def setKey(self, key):
        self.key = key

    def getLabelMap(self):
        return self.labelmap

    def getAlphaMap(self):
        return self.alphamap

    def getTexture(self):
        return self.texture

    def getCenters(self):
        return self.center

    def getRadius(self):
        return self.radius

    def getKey(self):
        return self.key

    def getIfColli(self):
        return self.ifcol

    def getCollisionList(self):
        return self.col_with

    def getUseLabel(self):
        return self.use_label