REVA-QCAV/src/utils/paste.py

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 = []
        for _ in range(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])
            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"