|
|
|
|
|
|
|
import math |
|
import random |
|
from random import randint |
|
|
|
import cv2 |
|
import numpy as np |
|
from PIL import Image, ImageDraw, ImageEnhance |
|
from skimage.measure import label as skimage_label |
|
|
|
from .raindrop import Raindrop, make_bezier |
|
|
|
|
|
def CheckCollision(DropList): |
|
"""This function handle the collision of the drops. |
|
|
|
:param DropList: list of raindrop class objects |
|
""" |
|
listFinalDrops = [] |
|
Checked_list = [] |
|
list_len = len(DropList) |
|
|
|
|
|
DropList.reverse() |
|
drop_key = 1 |
|
for drop in DropList: |
|
|
|
if drop.getKey() not in Checked_list: |
|
|
|
if drop.getIfColli(): |
|
|
|
collision_list = drop.getCollisionList() |
|
|
|
final_x = drop.getCenters()[0] * drop.getRadius() |
|
final_y = drop.getCenters()[1] * drop.getRadius() |
|
tmp_devide = drop.getRadius() |
|
final_R = drop.getRadius() * drop.getRadius() |
|
for col_id in collision_list: |
|
col_id = int(col_id) |
|
Checked_list.append(col_id) |
|
|
|
final_x += DropList[list_len - col_id].getRadius() * DropList[list_len - col_id].getCenters()[0] |
|
final_y += DropList[list_len - col_id].getRadius() * DropList[list_len - col_id].getCenters()[1] |
|
tmp_devide += DropList[list_len - col_id].getRadius() |
|
final_R += DropList[list_len - col_id].getRadius() * DropList[list_len - col_id].getRadius() |
|
final_x = int(round(final_x / tmp_devide)) |
|
final_y = int(round(final_y / tmp_devide)) |
|
final_R = int(round(math.sqrt(final_R))) |
|
|
|
newDrop = Raindrop(drop_key, (final_x, final_y), final_R) |
|
drop_key = drop_key + 1 |
|
listFinalDrops.append(newDrop) |
|
|
|
else: |
|
drop.setKey(drop_key) |
|
drop_key = drop_key + 1 |
|
listFinalDrops.append(drop) |
|
|
|
return listFinalDrops |
|
|
|
|
|
def generate_label(h, w, cfg): |
|
"""This function generate list of raindrop class objects and label map of |
|
this drops in the image. |
|
|
|
:param h: image height |
|
:param w: image width |
|
:param cfg: config with global constants |
|
:param shape: int from 0 to 2 defining raindrop shape type |
|
""" |
|
maxDrop = cfg['maxDrops'] |
|
minDrop = cfg['minDrops'] |
|
maxR = cfg['maxR'] |
|
minR = cfg['minR'] |
|
drop_num = randint(minDrop, maxDrop) |
|
imgh = h |
|
imgw = w |
|
|
|
ran_pos = [(int(random.random() * imgw), int(random.random() * imgh)) for _ in range(drop_num)] |
|
listRainDrops = [] |
|
listFinalDrops = [] |
|
for key, pos in enumerate(ran_pos): |
|
key = key + 1 |
|
radius = random.randint(minR, maxR) |
|
shape = random.randint(1, 1) |
|
drop = Raindrop(key, pos, radius, shape) |
|
listRainDrops.append(drop) |
|
|
|
label_map = np.zeros([h, w]) |
|
collisionNum = len(listRainDrops) |
|
listFinalDrops = list(listRainDrops) |
|
loop = 0 |
|
while collisionNum > 0: |
|
loop = loop + 1 |
|
listFinalDrops = list(listFinalDrops) |
|
collisionNum = len(listFinalDrops) |
|
label_map = np.zeros_like(label_map) |
|
|
|
for drop in listFinalDrops: |
|
|
|
(ix, iy) = drop.getCenters() |
|
radius = drop.getRadius() |
|
ROI_WL = 2 * radius |
|
ROI_WR = 2 * radius |
|
ROI_HU = 3 * radius |
|
ROI_HD = 2 * radius |
|
if (iy - 3 * radius) < 0: |
|
ROI_HU = iy |
|
if (iy + 2 * radius) > imgh: |
|
ROI_HD = imgh - iy |
|
if (ix - 2 * radius) < 0: |
|
ROI_WL = ix |
|
if (ix + 2 * radius) > imgw: |
|
ROI_WR = imgw - ix |
|
|
|
|
|
|
|
drop_label = drop.getLabelMap() |
|
|
|
if (label_map[iy, ix] > 0): |
|
col_ids = np.unique(label_map[iy - ROI_HU:iy + ROI_HD, ix - ROI_WL:ix + ROI_WR]) |
|
col_ids = col_ids[col_ids != 0] |
|
drop.setCollision(True, col_ids) |
|
label_map[iy - ROI_HU:iy + ROI_HD, |
|
ix - ROI_WL:ix + ROI_WR] = drop_label[3 * radius - ROI_HU:3 * radius + ROI_HD, 2 * radius - |
|
ROI_WL:2 * radius + ROI_WR] * drop.getKey() |
|
else: |
|
label_map[iy - ROI_HU:iy + ROI_HD, |
|
ix - ROI_WL:ix + ROI_WR] = drop_label[3 * radius - ROI_HU:3 * radius + ROI_HD, 2 * radius - |
|
ROI_WL:2 * radius + ROI_WR] * drop.getKey() |
|
|
|
collisionNum = collisionNum - 1 |
|
|
|
if collisionNum > 0: |
|
listFinalDrops = CheckCollision(listFinalDrops) |
|
return listFinalDrops, label_map |
|
|
|
|
|
def generateDrops(imagePath, cfg, listFinalDrops): |
|
"""Generate raindrops on the image. |
|
|
|
:param imagePath: path to the image on which you want to generate drops |
|
:param cfg: config with global constants |
|
:param listFinalDrops: final list of raindrop class objects after handling collisions |
|
:param label_map: general label map of all drops in the image |
|
""" |
|
ifReturnLabel = cfg['return_label'] |
|
edge_ratio = cfg['edge_darkratio'] |
|
|
|
PIL_bg_img = Image.open(imagePath).convert('RGB') |
|
bg_img = np.asarray(PIL_bg_img) |
|
label_map = np.zeros_like(bg_img)[:, :, 0] |
|
imgh, imgw, _ = bg_img.shape |
|
|
|
A = cfg['A'] |
|
B = cfg['B'] |
|
C = cfg['C'] |
|
D = cfg['D'] |
|
|
|
alpha_map = np.zeros_like(label_map).astype(np.float64) |
|
|
|
for drop in listFinalDrops: |
|
(ix, iy) = drop.getCenters() |
|
radius = drop.getRadius() |
|
ROI_WL = 2 * radius |
|
ROI_WR = 2 * radius |
|
ROI_HU = 3 * radius |
|
ROI_HD = 2 * radius |
|
if (iy - 3 * radius) < 0: |
|
ROI_HU = iy |
|
if (iy + 2 * radius) > imgh: |
|
ROI_HD = imgh - iy |
|
if (ix - 2 * radius) < 0: |
|
ROI_WL = ix |
|
if (ix + 2 * radius) > imgw: |
|
ROI_WR = imgw - ix |
|
|
|
drop_alpha = drop.getAlphaMap() |
|
alpha_map[iy - ROI_HU:iy + ROI_HD, |
|
ix - ROI_WL:ix + ROI_WR] += drop_alpha[3 * radius - ROI_HU:3 * radius + ROI_HD, |
|
2 * radius - ROI_WL:2 * radius + ROI_WR] |
|
|
|
alpha_map = alpha_map / np.max(alpha_map) * 255.0 |
|
|
|
PIL_bg_img = Image.open(imagePath) |
|
for idx, drop in enumerate(listFinalDrops): |
|
(ix, iy) = drop.getCenters() |
|
radius = drop.getRadius() |
|
ROIU = iy - 3 * radius |
|
ROID = iy + 2 * radius |
|
ROIL = ix - 2 * radius |
|
ROIR = ix + 2 * radius |
|
if (iy - 3 * radius) < 0: |
|
ROIU = 0 |
|
ROID = 5 * radius |
|
if (iy + 2 * radius) > imgh: |
|
ROIU = imgh - 5 * radius |
|
ROID = imgh |
|
if (ix - 2 * radius) < 0: |
|
ROIL = 0 |
|
ROIR = 4 * radius |
|
if (ix + 2 * radius) > imgw: |
|
ROIL = imgw - 4 * radius |
|
ROIR = imgw |
|
|
|
tmp_bg = bg_img[ROIU:ROID, ROIL:ROIR, :] |
|
try: |
|
drop.updateTexture(tmp_bg) |
|
except Exception: |
|
del listFinalDrops[idx] |
|
continue |
|
tmp_alpha_map = alpha_map[ROIU:ROID, ROIL:ROIR] |
|
|
|
output = drop.getTexture() |
|
tmp_output = np.asarray(output).astype(np.float)[:, :, -1] |
|
tmp_alpha_map = tmp_alpha_map * (tmp_output / 255) |
|
tmp_alpha_map = Image.fromarray(tmp_alpha_map.astype('uint8')) |
|
|
|
edge = ImageEnhance.Brightness(output) |
|
edge = edge.enhance(edge_ratio) |
|
|
|
PIL_bg_img.paste(edge, (ix - 2 * radius, iy - 3 * radius), output) |
|
PIL_bg_img.paste(output, (ix - 2 * radius, iy - 3 * radius), output) |
|
|
|
mask = np.zeros_like(bg_img) |
|
|
|
if len(listFinalDrops) > 0: |
|
|
|
for drop in listFinalDrops: |
|
if (drop.shape == 0): |
|
cv2.circle(mask, drop.center, drop.radius, (255, 255, 255), -1) |
|
if (drop.shape == 1): |
|
cv2.circle(mask, drop.center, drop.radius, (255, 255, 255), -1) |
|
cv2.ellipse(mask, drop.center, (drop.radius, int(1.3 * math.sqrt(3) * drop.radius)), 0, 180, 360, |
|
(255, 255, 255), -1) |
|
|
|
img = Image.fromarray(np.uint8(mask[:, :, 0]), 'L') |
|
|
|
for drop in listFinalDrops: |
|
if (drop.shape == 2): |
|
img = Image.fromarray(np.uint8(img), 'L') |
|
draw = ImageDraw.Draw(img) |
|
ts = [t / 100.0 for t in range(101)] |
|
xys = [(drop.radius * C[0] - 2 * drop.radius + drop.center[0], |
|
drop.radius * C[1] - 3 * drop.radius + drop.center[1]), |
|
(drop.radius * B[0] - 2 * drop.radius + drop.center[0], |
|
drop.radius * B[1] - 3 * drop.radius + drop.center[1]), |
|
(drop.radius * D[0] - 2 * drop.radius + drop.center[0], |
|
drop.radius * D[1] - 3 * drop.radius + drop.center[1])] |
|
bezier = make_bezier(xys) |
|
points = bezier(ts) |
|
|
|
xys = [(drop.radius * C[0] - 2 * drop.radius + drop.center[0], |
|
drop.radius * C[1] - 3 * drop.radius + drop.center[1]), |
|
(drop.radius * A[0] - 2 * drop.radius + drop.center[0], |
|
drop.radius * A[1] - 3 * drop.radius + drop.center[1]), |
|
(drop.radius * D[0] - 2 * drop.radius + drop.center[0], |
|
drop.radius * D[1] - 3 * drop.radius + drop.center[1])] |
|
bezier = make_bezier(xys) |
|
points.extend(bezier(ts)) |
|
draw.polygon(points, fill='white') |
|
mask = np.array(img) |
|
|
|
im_mask = Image.fromarray(mask.astype('uint8')) |
|
|
|
if ifReturnLabel: |
|
output_label = np.array(alpha_map) |
|
output_label.flags.writeable = True |
|
output_label[output_label > 0] = 1 |
|
output_label = Image.fromarray(output_label.astype('uint8')) |
|
return PIL_bg_img, output_label, im_mask |
|
|
|
return PIL_bg_img |
|
|
|
|
|
def generateDrops_np(img_np, cfg, listFinalDrops): |
|
"""Generate raindrops on the image. |
|
|
|
:param imgs: numpy imgs shape -> [B, H, W, C], type -> np.uint8 |
|
:param cfg: config with global constants |
|
:param listFinalDrops: final list of raindrop class objects after handling collisions |
|
:param label_map: general label map of all drops in the image |
|
""" |
|
ifReturnLabel = cfg['return_label'] |
|
edge_ratio = cfg['edge_darkratio'] |
|
|
|
|
|
|
|
|
|
bg_img = img_np |
|
label_map = np.zeros_like(bg_img)[:, :, 0] |
|
imgh, imgw, _ = bg_img.shape |
|
|
|
A = cfg['A'] |
|
B = cfg['B'] |
|
C = cfg['C'] |
|
D = cfg['D'] |
|
|
|
|
|
alpha_map = np.zeros_like(label_map).astype(np.float64) |
|
|
|
for drop in listFinalDrops: |
|
(ix, iy) = drop.getCenters() |
|
radius = drop.getRadius() |
|
ROI_WL = 2 * radius |
|
ROI_WR = 2 * radius |
|
ROI_HU = 3 * radius |
|
ROI_HD = 2 * radius |
|
if (iy - 3 * radius) < 0: |
|
ROI_HU = iy |
|
if (iy + 2 * radius) > imgh: |
|
ROI_HD = imgh - iy |
|
if (ix - 2 * radius) < 0: |
|
ROI_WL = ix |
|
if (ix + 2 * radius) > imgw: |
|
ROI_WR = imgw - ix |
|
|
|
drop_alpha = drop.getAlphaMap() |
|
|
|
alpha_map[iy - ROI_HU:iy + ROI_HD, |
|
ix - ROI_WL:ix + ROI_WR] += drop_alpha[3 * radius - ROI_HU:3 * radius + ROI_HD, |
|
2 * radius - ROI_WL:2 * radius + ROI_WR] |
|
|
|
alpha_map = alpha_map / np.max(alpha_map) * 255.0 |
|
|
|
PIL_bg_img = Image.fromarray(np.uint8(bg_img)).convert('RGB') |
|
|
|
for idx, drop in enumerate(listFinalDrops): |
|
(ix, iy) = drop.getCenters() |
|
radius = drop.getRadius() |
|
ROIU = iy - 3 * radius |
|
ROID = iy + 2 * radius |
|
ROIL = ix - 2 * radius |
|
ROIR = ix + 2 * radius |
|
if (iy - 3 * radius) < 0: |
|
ROIU = 0 |
|
ROID = 5 * radius |
|
if (iy + 2 * radius) > imgh: |
|
ROIU = imgh - 5 * radius |
|
ROID = imgh |
|
if (ix - 2 * radius) < 0: |
|
ROIL = 0 |
|
ROIR = 4 * radius |
|
if (ix + 2 * radius) > imgw: |
|
ROIL = imgw - 4 * radius |
|
ROIR = imgw |
|
|
|
tmp_bg = bg_img[ROIU:ROID, ROIL:ROIR] |
|
try: |
|
drop.updateTexture(tmp_bg) |
|
except Exception: |
|
del listFinalDrops[idx] |
|
continue |
|
tmp_alpha_map = alpha_map[ROIU:ROID, ROIL:ROIR] |
|
|
|
output = drop.getTexture() |
|
tmp_output = np.asarray(output).astype(np.float)[:, :, -1] |
|
tmp_alpha_map = tmp_alpha_map * (tmp_output / 255) |
|
tmp_alpha_map = Image.fromarray(tmp_alpha_map.astype('uint8')) |
|
|
|
edge = ImageEnhance.Brightness(output) |
|
edge = edge.enhance(edge_ratio) |
|
|
|
|
|
|
|
PIL_bg_img.paste(edge, (ROIL, ROIU), output) |
|
PIL_bg_img.paste(output, (ROIL, ROIU), output) |
|
|
|
|
|
|
|
mask = np.zeros_like(bg_img) |
|
|
|
if len(listFinalDrops) > 0: |
|
|
|
for drop in listFinalDrops: |
|
if (drop.shape == 0): |
|
cv2.circle(mask, drop.center, drop.radius, (255, 255, 255), -1) |
|
if (drop.shape == 1): |
|
cv2.circle(mask, drop.center, drop.radius, (255, 255, 255), -1) |
|
cv2.ellipse(mask, drop.center, (drop.radius, int(1.3 * math.sqrt(3) * drop.radius)), 0, 180, 360, |
|
(255, 255, 255), -1) |
|
|
|
img = Image.fromarray(np.uint8(mask[:, :, 0]), 'L') |
|
|
|
for drop in listFinalDrops: |
|
if (drop.shape == 2): |
|
img = Image.fromarray(np.uint8(img), 'L') |
|
draw = ImageDraw.Draw(img) |
|
ts = [t / 100.0 for t in range(101)] |
|
A0, A1 = drop.control_point['A'][0], drop.control_point['A'][1] |
|
B0, B1 = drop.control_point['B'][0], drop.control_point['B'][1] |
|
C0, C1 = drop.control_point['C'][0], drop.control_point['C'][1] |
|
D0, D1 = drop.control_point['D'][0], drop.control_point['D'][1] |
|
xys = [(drop.radius * C0 - 2 * drop.radius + drop.center[0], |
|
drop.radius * C1 - 3 * drop.radius + drop.center[1]), |
|
(drop.radius * B0 - 2 * drop.radius + drop.center[0], |
|
drop.radius * B1 - 3 * drop.radius + drop.center[1]), |
|
(drop.radius * D0 - 2 * drop.radius + drop.center[0], |
|
drop.radius * D1 - 3 * drop.radius + drop.center[1])] |
|
bezier = make_bezier(xys) |
|
points = bezier(ts) |
|
|
|
xys = [(drop.radius * C0 - 2 * drop.radius + drop.center[0], |
|
drop.radius * C1 - 3 * drop.radius + drop.center[1]), |
|
(drop.radius * A0 - 2 * drop.radius + drop.center[0], |
|
drop.radius * A1 - 3 * drop.radius + drop.center[1]), |
|
(drop.radius * D0 - 2 * drop.radius + drop.center[0], |
|
drop.radius * D1 - 3 * drop.radius + drop.center[1])] |
|
bezier = make_bezier(xys) |
|
points.extend(bezier(ts)) |
|
draw.polygon(points, fill='white') |
|
mask = np.array(img) |
|
|
|
im_mask = Image.fromarray(mask.astype('uint8')) |
|
|
|
if ifReturnLabel: |
|
output_label = np.array(alpha_map) |
|
output_label.flags.writeable = True |
|
output_label[output_label > 0] = 1 |
|
output_label = Image.fromarray(output_label.astype('uint8')) |
|
return PIL_bg_img, output_label, im_mask |
|
|
|
return PIL_bg_img |
|
|