Spaces:
Running
on
Zero
Running
on
Zero
# -*- coding: utf-8 -*- | |
# Copyright (c) Alibaba, Inc. and its affiliates. | |
from abc import ABCMeta | |
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 dict_to_yaml | |
class CannyAnnotator(BaseAnnotator, metaclass=ABCMeta): | |
para_dict = {} | |
def __init__(self, cfg, logger=None): | |
super().__init__(cfg, logger=logger) | |
self.low_threshold = cfg.get('LOW_THRESHOLD', 100) | |
self.high_threshold = cfg.get('HIGH_THRESHOLD', 200) | |
self.random_cfg = cfg.get('RANDOM_CFG', None) | |
def forward(self, image): | |
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.' | |
assert len(image.shape) < 4 | |
if self.random_cfg is None: | |
image = cv2.Canny(image, self.low_threshold, self.high_threshold) | |
else: | |
proba = self.random_cfg.get('PROBA', 1.0) | |
if np.random.random() < proba: | |
min_low_threshold = self.random_cfg.get( | |
'MIN_LOW_THRESHOLD', 50) | |
max_low_threshold = self.random_cfg.get( | |
'MAX_LOW_THRESHOLD', 100) | |
min_high_threshold = self.random_cfg.get( | |
'MIN_HIGH_THRESHOLD', 200) | |
max_high_threshold = self.random_cfg.get( | |
'MAX_HIGH_THRESHOLD', 350) | |
low_th = np.random.randint(min_low_threshold, | |
max_low_threshold) | |
high_th = np.random.randint(min_high_threshold, | |
max_high_threshold) | |
else: | |
low_th, high_th = self.low_threshold, self.high_threshold | |
image = cv2.Canny(image, low_th, high_th) | |
return image[..., None].repeat(3, 2) | |
def get_config_template(): | |
return dict_to_yaml('ANNOTATORS', | |
__class__.__name__, | |
CannyAnnotator.para_dict, | |
set_name=True) | |