# -*- coding: utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. import random from abc import ABCMeta from enum import Enum import cv2 import numpy as np import torch from PIL import Image from scepter.modules.annotator.base_annotator import BaseAnnotator from scepter.modules.annotator.registry import ANNOTATORS from scepter.modules.utils.config import Config, dict_to_yaml def invert_image(im): im_arr = np.array(im) mask_1 = im_arr == 0 mask_2 = im_arr == 255 im_arr[mask_1] = 255 im_arr[mask_2] = 0 new_im = im_arr return new_im class DrawMethod(Enum): LINE = 'line' CIRCLE = 'circle' SQUARE = 'square' def make_random_irregular_mask(shape, max_angle=4, max_len=60, max_width=20, min_times=0, max_times=10, draw_method=DrawMethod.LINE): draw_method = DrawMethod(draw_method) height, width = shape mask = np.zeros((height, width), np.float32) times = np.random.randint(min_times, max_times + 1) for i in range(times): start_x = np.random.randint(width) start_y = np.random.randint(height) for j in range(1 + np.random.randint(5)): angle = 0.01 + np.random.randint(max_angle) if i % 2 == 0: angle = 2 * 3.1415926 - angle length = 10 + np.random.randint(max_len) brush_w = 5 + np.random.randint(max_width) end_x = np.clip( (start_x + length * np.sin(angle)).astype(np.int32), 0, width) end_y = np.clip( (start_y + length * np.cos(angle)).astype(np.int32), 0, height) if draw_method == DrawMethod.LINE: cv2.line(mask, (start_x, start_y), (end_x, end_y), 1.0, brush_w) elif draw_method == DrawMethod.CIRCLE: cv2.circle(mask, (start_x, start_y), radius=brush_w, color=1., thickness=-1) elif draw_method == DrawMethod.SQUARE: radius = brush_w // 2 mask[start_y - radius:start_y + radius, start_x - radius:start_x + radius] = 1 start_x, start_y = end_x, end_y return mask[None, ...] class RandomIrregularMaskGenerator: def __init__(self, max_angle=4, max_len=60, max_width=20, min_times=0, max_times=10, ramp_kwargs=None, draw_method=DrawMethod.LINE): self.max_angle = max_angle self.max_len = max_len self.max_width = max_width self.min_times = min_times self.max_times = max_times self.draw_method = draw_method self.ramp = None # self.ramp = LinearRamp(**ramp_kwargs) if ramp_kwargs is not None else None def __call__(self, img, iter_i=None, raw_image=None): coef = self.ramp(iter_i) if (self.ramp is not None) and ( iter_i is not None) else 1 cur_max_len = int(max(1, self.max_len * coef)) cur_max_width = int(max(1, self.max_width * coef)) cur_max_times = int(self.min_times + 1 + (self.max_times - self.min_times) * coef) return make_random_irregular_mask(img.shape[1:], max_angle=self.max_angle, max_len=cur_max_len, max_width=cur_max_width, min_times=self.min_times, max_times=cur_max_times, draw_method=self.draw_method) def make_random_rectangle_mask(shape, margin=10, bbox_min_size=30, bbox_max_size=100, min_times=0, max_times=3): height, width = shape mask = np.zeros((height, width), np.float32) bbox_max_size = min(bbox_max_size, height - margin * 2, width - margin * 2) times = np.random.randint(min_times, max_times + 1) for i in range(times): box_width = np.random.randint(bbox_min_size, bbox_max_size) box_height = np.random.randint(bbox_min_size, bbox_max_size) start_x = np.random.randint(margin, width - margin - box_width + 1) start_y = np.random.randint(margin, height - margin - box_height + 1) mask[start_y:start_y + box_height, start_x:start_x + box_width] = 1 return mask[None, ...] class RandomRectangleMaskGenerator: def __init__(self, margin=10, bbox_min_size=30, bbox_max_size=100, min_times=0, max_times=3, ramp_kwargs=None): self.margin = margin self.bbox_min_size = bbox_min_size self.bbox_max_size = bbox_max_size self.min_times = min_times self.max_times = max_times self.ramp = None # self.ramp = LinearRamp(**ramp_kwargs) if ramp_kwargs is not None else None def __call__(self, img, iter_i=None, raw_image=None): coef = self.ramp(iter_i) if (self.ramp is not None) and ( iter_i is not None) else 1 cur_bbox_max_size = int(self.bbox_min_size + 1 + (self.bbox_max_size - self.bbox_min_size) * coef) cur_max_times = int(self.min_times + (self.max_times - self.min_times) * coef) return make_random_rectangle_mask(img.shape[1:], margin=self.margin, bbox_min_size=self.bbox_min_size, bbox_max_size=cur_bbox_max_size, min_times=self.min_times, max_times=cur_max_times) class MixedMaskGenerator: def __init__(self, irregular_proba=1 / 3, irregular_kwargs=None, box_proba=1 / 3, box_kwargs=None, invert_proba=0): self.probas = [] self.gens = [] if irregular_proba > 0: self.probas.append(irregular_proba) if irregular_kwargs is None: irregular_kwargs = {} else: irregular_kwargs = dict(irregular_kwargs) irregular_kwargs['draw_method'] = DrawMethod.LINE self.gens.append(RandomIrregularMaskGenerator(**irregular_kwargs)) if box_proba > 0: self.probas.append(box_proba) if box_kwargs is None: box_kwargs = {} self.gens.append(RandomRectangleMaskGenerator(**box_kwargs)) self.probas = np.array(self.probas, dtype='float32') self.probas /= self.probas.sum() self.invert_proba = invert_proba def __call__(self, img, iter_i=None, raw_image=None): kind = np.random.choice(len(self.probas), p=self.probas) gen = self.gens[kind] result = gen(img, iter_i=iter_i, raw_image=raw_image) if self.invert_proba > 0 and random.random() < self.invert_proba: result = 1 - result return result @ANNOTATORS.register_class() class InpaintingAnnotator(BaseAnnotator, metaclass=ABCMeta): para_dict = {} def __init__(self, cfg, logger=None): super().__init__(cfg, logger=logger) self.mask_cfg = cfg.get( 'MASK_CFG', { 'irregular_proba': 0.5, 'irregular_kwargs': { 'min_times': 4, 'max_times': 10, 'max_width': 100, 'max_angle': 4, 'max_len': 200 }, 'box_proba': 0.5, 'box_kwargs': { 'margin': 0, 'bbox_min_size': 30, 'bbox_max_size': 150, 'max_times': 5, 'min_times': 1 } }) self.mask_cfg = Config.get_dict(self.mask_cfg) if not isinstance( self.mask_cfg, dict) else self.mask_cfg self.mask_generator = MixedMaskGenerator(**self.mask_cfg) self.return_mask = cfg.get('RETURN_MASK', False) self.return_invert = cfg.get('RETURN_INVERT', True) self.mask_color = cfg.get('MASK_COLOR', 0) def forward(self, image, mask=None, return_mask=None, return_invert=None, mask_color=None): return_mask = return_mask if return_mask is not None else self.return_mask return_invert = return_invert if return_invert is not None else self.return_invert mask_color = mask_color if mask_color is not None else self.mask_color if isinstance(image, Image.Image): image = np.array(image) elif isinstance(image, torch.Tensor): image = image.detach().cpu().numpy() elif isinstance(image, np.ndarray): image = image.copy() else: raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.' if mask is not None: if mask_color: image[np.array(mask) == 255] = mask_color else: image[np.array(mask) == 255] = 0 else: img = np.transpose(image, (2, 0, 1)) mask = self.mask_generator(img) mask = (np.transpose(mask, (1, 2, 0)).squeeze(-1) * 255).astype(np.uint8) if return_invert: mask = invert_image(mask) colored_mask = np.zeros_like(image) if mask_color: colored_mask[:] = mask_color image = np.where(mask[:, :, np.newaxis] == 255, colored_mask, image) if return_mask: ret_data = {'image': np.array(image), 'mask': np.array(mask)} else: ret_data = np.array(image) return ret_data @staticmethod def get_config_template(): return dict_to_yaml('ANNOTATORS', __class__.__name__, InpaintingAnnotator.para_dict, set_name=True)