black out (mask) polygons from an image with albumentations

blackout_polygons_aug.py

import numpy as np
import albumentations as A
import cv2
from albumentations.pytorch import ToTensorV2
from albumentations.augmentations import functional as A_F
from albumentations.core.transforms_interface import ImageOnlyTransform

class BlackOutWithPolygonMasks(ImageOnlyTransform):
    def __init__(self, always_apply=False, p=1.0):
        super(BlackOutWithPolygonMasks, self).__init__(always_apply, p)

    @property
    def targets_as_params(self):
        return ["image", "polygon_masks", "start_coordinates"]

    def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, Any]:
        return params

    def apply(self, img, **params):
        masks = params.get('polygon_masks')
        start_coordinates = params.get('start_coordinates')
        if masks is None or len(masks) == 0:
            return img

        if start_coordinates is not None:
            masks = list(map(lambda x: x[:, 0] - start_coordinates, masks))

        # Create a black mask with the same shape as the image
        black_mask = np.ones_like(img) * 255

        for mask in masks:
            # Draw the polygon on the black mask
            cv2.fillPoly(black_mask, [mask], (0, 0, 0))

        # Use the black mask to black out the image where the polygon is
        img_blackout = cv2.bitwise_and(img, black_mask)

        return img_blackout
 """
 I use it for masking speech bubbles from comic book panels...
# sample usage... 
masks = panel_info['panel_speech_masks']
transform = A.Compose([
    BlackOutWithPolygonMasks(always_apply=True, p=1.0),
    self.transform_panel
])
panel_img = transform(image=panel_img,
                      polygon_masks=masks,
                      start_coordinates=panel_info['panel_start_coordinates'])['image']
 """