diff --git a/modules/image/image_processing/enlightengan/enlighten_inference/__init__.py b/modules/image/image_processing/enlightengan/enlighten_inference/__init__.py deleted file mode 100755 index 5736aafcbcd81bdfdfe284a177005f98495354d3..0000000000000000000000000000000000000000 --- a/modules/image/image_processing/enlightengan/enlighten_inference/__init__.py +++ /dev/null @@ -1,37 +0,0 @@ -import os -from typing import Union - -import numpy as np -from onnxruntime import InferenceSession - - -def get_relative_path(root, *args): - return os.path.join(os.path.dirname(root), *args) - - -class EnlightenOnnxModel: - def __init__(self, model: Union[bytes, str, None] = None): - self.graph = InferenceSession(model or get_relative_path(__file__, 'enlighten.onnx')) - - def __repr__(self): - return ' {}'.format(id(self)) - - def _pad(self, img): - h, w, _ = img.shape - block_size = 16 - min_height = (h // block_size + 1) * block_size - min_width = (w // block_size + 1) * block_size - img = np.pad(img, ((0, min_height - h), (0, min_width - w), (0, 0)), mode='constant', constant_values=0) - return img, (h, w) - - def _preprocess(self, img): - if len(img.shape) != 3: - raise ValueError('Incorrect shape: expected 3, got {}'.format(img.shape)) - return np.expand_dims(np.transpose(img, (2, 0, 1)).astype(np.float32) / 255., 0) - - def predict(self, img): - padded, (h, w) = self._pad(img) - image_numpy, = self.graph.run(['output'], {'input': self._preprocess(padded)}) - image_numpy = (np.transpose(image_numpy[0], (1, 2, 0)) + 1) / 2.0 * 255.0 - image_numpy = np.clip(image_numpy, 0, 255) - return image_numpy.astype('uint8')[:h, :w, :] diff --git a/modules/image/image_processing/enlightengan/enlighten_inference/pd_model/x2paddle_code.py b/modules/image/image_processing/enlightengan/enlighten_inference/pd_model/x2paddle_code.py index 51e7a5e61e669edc72a174cb9db06c95a3e24029..d211efac274f7d6be42ee8a765726526d9e51888 100755 --- a/modules/image/image_processing/enlightengan/enlighten_inference/pd_model/x2paddle_code.py +++ b/modules/image/image_processing/enlightengan/enlighten_inference/pd_model/x2paddle_code.py @@ -1,6 +1,5 @@ import paddle import math -from x2paddle.op_mapper.onnx2paddle import onnx_custom_layer as x2paddle_nn class ONNXModel(paddle.nn.Layer): diff --git a/modules/image/image_processing/enlightengan/module.py b/modules/image/image_processing/enlightengan/module.py index eacfaf06e1bbf0ab8c10a3c162512687ec56fe22..30749182bde2f984fb1b73624b01a56010fa4e1d 100644 --- a/modules/image/image_processing/enlightengan/module.py +++ b/modules/image/image_processing/enlightengan/module.py @@ -21,7 +21,6 @@ from paddlehub.module.module import moduleinfo, runnable, serving import numpy as np import cv2 -from .enlighten_inference import EnlightenOnnxModel from .enlighten_inference.pd_model.x2paddle_code import ONNXModel from .util import base64_to_cv2 @@ -36,19 +35,19 @@ class EnlightenGAN: self.model.set_dict(params, use_structured_name=True) def enlightening(self, - images=None, - paths=None, - output_dir='./enlightening_result/', - use_gpu=False, - visualization=True): + images:list=None, + paths:list=None, + output_dir:str='./enlightening_result/', + use_gpu:bool=False, + visualization:bool=True): ''' enlighten images in the low-light scene. images (list[numpy.ndarray]): data of images, shape of each is [H, W, C], color space must be BGR(read by cv2). paths (list[str]): paths to images - output_dir: the dir to save the results - use_gpu: if True, use gpu to perform the computation, otherwise cpu. - visualization: if True, save results in output_dir. + output_dir (str): the dir to save the results + use_gpu (bool): if True, use gpu to perform the computation, otherwise cpu. + visualization (bool): if True, save results in output_dir. ''' results = [] paddle.disable_static() diff --git a/modules/image/image_processing/enlightengan/requirements.txt b/modules/image/image_processing/enlightengan/requirements.txt deleted file mode 100644 index 8e095c08f78bd8d5d02bba769dfdb65b54e06c43..0000000000000000000000000000000000000000 --- a/modules/image/image_processing/enlightengan/requirements.txt +++ /dev/null @@ -1,2 +0,0 @@ -onnxruntime -x2paddle