136 lines
4.5 KiB
Python
136 lines
4.5 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import math
|
|
import random
|
|
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 Config, dict_to_yaml
|
|
|
|
|
|
def gaussian_noise_op(im, v):
|
|
from basicsr.data.degradations import random_add_gaussian_noise
|
|
noise_level = v.get('noise_level', [10, 20])
|
|
out = random_add_gaussian_noise(
|
|
im,
|
|
sigma_range=noise_level,
|
|
clip=True,
|
|
rounds=False,
|
|
gray_prob=0.4,
|
|
)
|
|
out = np.clip(out, 0.0, 1.0)
|
|
return out
|
|
|
|
|
|
def resize_op(im, v):
|
|
scale = v.get('scale', [0.5, 0.8])
|
|
h, w = im.shape[:2]
|
|
scale = random.uniform(scale[0], scale[1])
|
|
h_, w_ = int(h * scale), int(w * scale)
|
|
mode = v.get('mode', 'nearest')
|
|
if mode == 'nearest':
|
|
interpolation = cv2.INTER_NEAREST
|
|
elif mode == 'bilinear':
|
|
interpolation = cv2.INTER_LINEAR
|
|
elif mode == 'bicubic':
|
|
interpolation = cv2.INTER_CUBIC
|
|
else:
|
|
interpolation = cv2.INTER_NEAREST
|
|
im = cv2.resize(im, (w_, h_), interpolation=interpolation)
|
|
out = cv2.resize(im, (w, h), interpolation=interpolation)
|
|
out = np.clip(out, 0.0, 1.0)
|
|
return out
|
|
|
|
|
|
def jpeg_op(im, v):
|
|
from basicsr.data.degradations import add_jpg_compression
|
|
jpeg_level = v.get('jpeg_level', [50, 75])
|
|
v = int(random.uniform(jpeg_level[0], jpeg_level[1]))
|
|
out = add_jpg_compression(im, v)
|
|
out = np.clip(out, 0.0, 1.0)
|
|
return out
|
|
|
|
|
|
def gaussian_blur_op(im, v):
|
|
from basicsr.data.degradations import random_mixed_kernels
|
|
kernel_range = v.get('kernel_size', [7, 9])
|
|
kernel_size = random.choice(kernel_range)
|
|
kernel_size = min(int(kernel_size) // 2 * 2 + 1, 21)
|
|
blur_sigma = v.get('sigma', [0.9, 1.0])
|
|
kernel = random_mixed_kernels(
|
|
('iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso',
|
|
'plateau_aniso'), (0.45, 0.25, 0.12, 0.03, 0.12, 0.03),
|
|
kernel_size,
|
|
blur_sigma,
|
|
blur_sigma, [-math.pi, math.pi], [0.5, 2.0], [1, 1.5],
|
|
noise_range=None)
|
|
|
|
pad_size = (21 - kernel_size) // 2
|
|
kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))
|
|
out = cv2.filter2D(im, -1, kernel)
|
|
out = np.clip(out, 0.0, 1.0)
|
|
return out
|
|
|
|
|
|
@ANNOTATORS.register_class()
|
|
class DegradationAnnotator(BaseAnnotator, metaclass=ABCMeta):
|
|
para_dict = {}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super().__init__(cfg, logger=logger)
|
|
self.params = cfg.get('PARAMS', {
|
|
'gaussian_noise': {},
|
|
'resize': {},
|
|
'jpeg': {},
|
|
'gaussian_blur': {},
|
|
})
|
|
if not isinstance(self.params, dict):
|
|
self.params = Config.get_dict(self.params)
|
|
self.random_degradation = cfg.get('RANDOM_DEGRADATION', False)
|
|
|
|
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.'
|
|
if np.max(image) > 1.0:
|
|
image = (image / 255.).astype(np.float32)
|
|
|
|
degradation_list = list(self.params.keys())
|
|
if self.random_degradation:
|
|
random.shuffle(degradation_list)
|
|
|
|
for degradation_type in degradation_list:
|
|
if degradation_type == 'gaussian_noise':
|
|
image = gaussian_noise_op(image, self.params[degradation_type])
|
|
elif degradation_type == 'resize':
|
|
image = resize_op(image, self.params[degradation_type])
|
|
elif degradation_type == 'jpeg':
|
|
image = jpeg_op(image, self.params[degradation_type])
|
|
elif degradation_type == 'gaussian_blur':
|
|
image = gaussian_blur_op(image, self.params[degradation_type])
|
|
else:
|
|
raise NotImplementedError(
|
|
f'ERROR: degradation_type: {degradation_type} is invalid.')
|
|
image = (image * 255.0).astype(np.uint8)
|
|
|
|
assert len(image.shape) < 4
|
|
return image
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('ANNOTATORS',
|
|
__class__.__name__,
|
|
DegradationAnnotator.para_dict,
|
|
set_name=True)
|