mirror of
https://github.com/Laurent2916/REVA-QCAV.git
synced 2024-11-14 09:08:18 +00:00
e089bcc7c4
Former-commit-id: 920a2df1e54e3dfc22cf62fe9c892bf945af9b63 [formerly f0c2daaa4d0094fafeddce9c07c6ce020c1704d6] Former-commit-id: 7e74f04ed0cd0d31612f23a44b9832b659bf5408
133 lines
3.7 KiB
Python
133 lines
3.7 KiB
Python
import os
|
|
import random as rd
|
|
|
|
import albumentations as A
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
class RandomPaste(A.DualTransform):
|
|
"""Paste an object on a background.
|
|
|
|
Args:
|
|
TODO
|
|
|
|
Targets:
|
|
image, mask
|
|
|
|
Image types:
|
|
uint8
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
nb,
|
|
path_paste_img_dir,
|
|
path_paste_mask_dir,
|
|
scale_range=(0.1, 0.2),
|
|
always_apply=True,
|
|
p=1.0,
|
|
):
|
|
super().__init__(always_apply, p)
|
|
self.path_paste_img_dir = path_paste_img_dir
|
|
self.path_paste_mask_dir = path_paste_mask_dir
|
|
self.scale_range = scale_range
|
|
self.nb = nb
|
|
|
|
@property
|
|
def targets_as_params(self):
|
|
return ["image"]
|
|
|
|
def apply(self, img, positions, paste_img, paste_mask, **params):
|
|
# convert img to Image, needed for `paste` function
|
|
img = Image.fromarray(img)
|
|
|
|
for pos in positions:
|
|
img.paste(paste_img, pos, paste_mask)
|
|
|
|
return np.asarray(img.convert("RGB"))
|
|
|
|
def apply_to_mask(self, mask, positions, paste_mask, **params):
|
|
# convert mask to Image, needed for `paste` function
|
|
mask = Image.fromarray(mask)
|
|
|
|
for pos in positions:
|
|
mask.paste(paste_mask, pos, paste_mask)
|
|
|
|
return np.asarray(mask.convert("L"))
|
|
|
|
def get_params_dependent_on_targets(self, params):
|
|
# choose a random image inside the image folder
|
|
filename = rd.choice(os.listdir(self.path_paste_img_dir))
|
|
|
|
# load the "paste" image
|
|
paste_img = Image.open(
|
|
os.path.join(
|
|
self.path_paste_img_dir,
|
|
filename,
|
|
)
|
|
).convert("RGBA")
|
|
|
|
# load its respective mask
|
|
paste_mask = Image.open(
|
|
os.path.join(
|
|
self.path_paste_mask_dir,
|
|
filename,
|
|
)
|
|
).convert("LA")
|
|
|
|
# load the target image
|
|
target_img = params["image"]
|
|
target_shape = np.array(target_img.shape[:2], dtype=np.uint)
|
|
paste_shape = np.array(paste_img.size, dtype=np.uint)
|
|
|
|
# compute the minimum scaling to fit inside target image
|
|
min_scale = np.min(target_shape / paste_shape)
|
|
|
|
# randomize the relative scaling
|
|
scale = rd.uniform(*self.scale_range)
|
|
|
|
# rotate the image and its mask
|
|
angle = rd.uniform(0, 360)
|
|
paste_img = paste_img.rotate(angle, expand=True)
|
|
paste_mask = paste_mask.rotate(angle, expand=True)
|
|
|
|
# scale the "paste" image and its mask
|
|
paste_img = paste_img.resize(
|
|
tuple((paste_shape * min_scale * scale).astype(np.uint)),
|
|
resample=Image.Resampling.LANCZOS,
|
|
)
|
|
paste_mask = paste_mask.resize(
|
|
tuple((paste_shape * min_scale * scale).astype(np.uint)),
|
|
resample=Image.Resampling.LANCZOS,
|
|
)
|
|
|
|
# update paste_shape after scaling
|
|
paste_shape = np.array(paste_img.size, dtype=np.uint)
|
|
|
|
# generate some positions
|
|
positions = []
|
|
while len(positions) <= rd.randint(1, self.nb):
|
|
x = rd.randint(0, target_shape[0] - paste_shape[0])
|
|
y = rd.randint(0, target_shape[1] - paste_shape[1])
|
|
|
|
# check for overlapping
|
|
for xo, yo in positions:
|
|
if (x <= xo + paste_shape[0]) and (y <= yo + paste_shape[1]):
|
|
continue
|
|
|
|
positions.append((x, y))
|
|
|
|
params.update(
|
|
{
|
|
"positions": positions,
|
|
"paste_img": paste_img,
|
|
"paste_mask": paste_mask,
|
|
}
|
|
)
|
|
|
|
return params
|
|
|
|
def get_transform_init_args_names(self):
|
|
return "scale_range", "path_paste_img_dir", "path_paste_mask_dir"
|