From fa8eabcf2f99e0277b6367e016ea11e0fbd0a0c5 Mon Sep 17 00:00:00 2001 From: gaotingquan Date: Wed, 9 Mar 2022 02:07:04 +0000 Subject: [PATCH] add preprocess.py --- .../processor/data_processor/preprocess.py | 268 ++++++++++++++++++ 1 file changed, 268 insertions(+) create mode 100644 deploy/python/ppshitu_v2/processor/data_processor/preprocess.py diff --git a/deploy/python/ppshitu_v2/processor/data_processor/preprocess.py b/deploy/python/ppshitu_v2/processor/data_processor/preprocess.py new file mode 100644 index 00000000..f4a202fb --- /dev/null +++ b/deploy/python/ppshitu_v2/processor/data_processor/preprocess.py @@ -0,0 +1,268 @@ +from functools import partial +import six +import math +import random +import cv2 +import numpy as np +import importlib +from PIL import Image + +from utils import logger + + +class PreProcesser(object): + def __init__(self, config): + """Image PreProcesser + + Args: + config (list): A list consisting of Dict object that describe an image processer operator. + """ + super().__init__() + self.ops = self.create_ops(config) + + def create_ops(self, config): + if not isinstance(config, list): + msg = "The preprocess config should be a list consisting of Dict object." + logger.error(msg) + raise Exception(msg) + mod = importlib.import_module(__name__) + ops = [] + for op_config in config: + name = list(op_config)[0] + param = {} if op_config[name] is None else op_config[name] + op = getattr(mod, name)(**param) + ops.append(op) + return ops + + def __call__(self, img, img_info=None): + if img_info: + for op in self.ops: + img, img_info = op(img, img_info) + return img, img_info + else: + for op in self.ops: + img = op(img) + return img + + +class DecodeImage(object): + """ decode image """ + + def __init__(self, to_rgb=True, to_np=False, channel_first=False): + self.to_rgb = to_rgb + self.to_np = to_np # to numpy + self.channel_first = channel_first # only enabled when to_np is True + + def __call__(self, img, img_info=None): + if six.PY2: + assert type(img) is str and len( + img) > 0, "invalid input 'img' in DecodeImage" + else: + assert type(img) is bytes and len( + img) > 0, "invalid input 'img' in DecodeImage" + data = np.frombuffer(img, dtype='uint8') + img = cv2.imdecode(data, 1) + if self.to_rgb: + assert img.shape[2] == 3, 'invalid shape of image[%s]' % ( + img.shape) + img = img[:, :, ::-1] + + if self.channel_first: + img = img.transpose((2, 0, 1)) + + if img_info: + img_info["im_shape"] = np.array(img.shape[:2], dtype=np.float32) + img_info["scale_factor"] = np.array([1., 1.], dtype=np.float32) + return img, img_info + else: + return img + + +class UnifiedResize(object): + def __init__(self, interpolation=None, backend="cv2"): + _cv2_interp_from_str = { + 'nearest': cv2.INTER_NEAREST, + 'bilinear': cv2.INTER_LINEAR, + 'area': cv2.INTER_AREA, + 'bicubic': cv2.INTER_CUBIC, + 'lanczos': cv2.INTER_LANCZOS4 + } + _pil_interp_from_str = { + 'nearest': Image.NEAREST, + 'bilinear': Image.BILINEAR, + 'bicubic': Image.BICUBIC, + 'box': Image.BOX, + 'lanczos': Image.LANCZOS, + 'hamming': Image.HAMMING + } + + def _pil_resize(src, size, resample): + pil_img = Image.fromarray(src) + pil_img = pil_img.resize(size, resample) + return np.asarray(pil_img) + + if backend.lower() == "cv2": + if isinstance(interpolation, str): + interpolation = _cv2_interp_from_str[interpolation.lower()] + # compatible with opencv < version 4.4.0 + elif interpolation is None: + interpolation = cv2.INTER_LINEAR + self.resize_func = partial(cv2.resize, interpolation=interpolation) + elif backend.lower() == "pil": + if isinstance(interpolation, str): + interpolation = _pil_interp_from_str[interpolation.lower()] + self.resize_func = partial(_pil_resize, resample=interpolation) + else: + logger.warning( + f"The backend of Resize only support \"cv2\" or \"PIL\". \"f{backend}\" is unavailable. Use \"cv2\" instead." + ) + self.resize_func = cv2.resize + + def __call__(self, src, size): + return self.resize_func(src, size) + + +class ResizeImage(object): + """ resize image """ + + def __init__(self, + size=None, + resize_short=None, + interpolation=None, + backend="cv2"): + if resize_short is not None and resize_short > 0: + self.resize_short = resize_short + self.w = None + self.h = None + elif size is not None: + self.resize_short = None + self.w = size if type(size) is int else size[0] + self.h = size if type(size) is int else size[1] + else: + raise Exception("invalid params for ReisizeImage for '\ + 'both 'size' and 'resize_short' are None") + + self._resize_func = UnifiedResize( + interpolation=interpolation, backend=backend) + + def __call__(self, img, img_info=None): + img_h, img_w = img.shape[:2] + if self.resize_short is not None: + percent = float(self.resize_short) / min(img_w, img_h) + w = int(round(img_w * percent)) + h = int(round(img_h * percent)) + else: + w = self.w + h = self.h + img = self._resize_func(img, (w, h)) + if img_info: + img_info["input_shape"] = img.shape[:2] + img_info["scale_factor"] = np.array( + [img.shape[0] / img_h, img.shape[1] / img_w]).astype("float32") + return img, img_info + else: + return img + + +class CropImage(object): + """ crop image """ + + def __init__(self, size): + if type(size) is int: + self.size = (size, size) + else: + self.size = size # (h, w) + + def __call__(self, img, img_info=None): + w, h = self.size + img_h, img_w = img.shape[:2] + + if img_h < h or img_w < w: + raise Exception( + f"The size({h}, {w}) of CropImage must be greater than size({img_h}, {img_w}) of image. Please check image original size and size of ResizeImage if used." + ) + + w_start = (img_w - w) // 2 + h_start = (img_h - h) // 2 + + w_end = w_start + w + h_end = h_start + h + img = img[h_start:h_end, w_start:w_end, :] + if img_info: + img_info["input_shape"] = img.shape[:2] + # TODO(gaotingquan): im_shape is needed to update? + img_info["im_shape"] = np.array(img.shape[:2], dtype=np.float32) + return img, img_info + else: + return img + + +class NormalizeImage(object): + """ normalize image such as substract mean, divide std + """ + + def __init__(self, + scale=None, + mean=None, + std=None, + order='chw', + output_fp16=False, + channel_num=3): + if isinstance(scale, str): + scale = eval(scale) + assert channel_num in [ + 3, 4 + ], "channel number of input image should be set to 3 or 4." + self.channel_num = channel_num + self.output_dtype = 'float16' if output_fp16 else 'float32' + self.scale = np.float32(scale if scale is not None else 1.0 / 255.0) + self.order = order + mean = mean if mean is not None else [0.485, 0.456, 0.406] + std = std if std is not None else [0.229, 0.224, 0.225] + + shape = (3, 1, 1) if self.order == 'chw' else (1, 1, 3) + self.mean = np.array(mean).reshape(shape).astype('float32') + self.std = np.array(std).reshape(shape).astype('float32') + + def __call__(self, img, img_info=None): + if isinstance(img, Image.Image): + img = np.array(img) + + assert isinstance(img, + np.ndarray), "invalid input 'img' in NormalizeImage" + + img = (img.astype('float32') * self.scale - self.mean) / self.std + + if self.channel_num == 4: + img_h = img.shape[1] if self.order == 'chw' else img.shape[0] + img_w = img.shape[2] if self.order == 'chw' else img.shape[1] + pad_zeros = np.zeros( + (1, img_h, img_w)) if self.order == 'chw' else np.zeros( + (img_h, img_w, 1)) + img = (np.concatenate( + (img, pad_zeros), axis=0) + if self.order == 'chw' else np.concatenate( + (img, pad_zeros), axis=2)) + img = img.astype(self.output_dtype) + if img_info: + return img, img_info + else: + return img + + +class ToCHWImage(object): + """ convert hwc image to chw image + """ + + def __init__(self): + pass + + def __call__(self, img, img_info=None): + if isinstance(img, Image.Image): + img = np.array(img) + + img = img.transpose((2, 0, 1)) + if img_info: + return img, img_info + else: + return img -- GitLab