module.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145
  1. # Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import absolute_import
  15. from __future__ import division
  16. from __future__ import print_function
  17. import os
  18. import sys
  19. sys.path.insert(0, ".")
  20. import copy
  21. import time
  22. import paddlehub
  23. from paddlehub.common.logger import logger
  24. from paddlehub.module.module import moduleinfo, runnable, serving
  25. import cv2
  26. import numpy as np
  27. import paddlehub as hub
  28. from tools.infer.utility import base64_to_cv2
  29. from ppstructure.kie.predict_kie_token_ser import SerPredictor
  30. from ppstructure.utility import parse_args
  31. from deploy.hubserving.kie_ser.params import read_params
  32. @moduleinfo(
  33. name="kie_ser",
  34. version="1.0.0",
  35. summary="kie ser service",
  36. author="paddle-dev",
  37. author_email="paddle-dev@baidu.com",
  38. type="cv/KIE_SER")
  39. class KIESer(hub.Module):
  40. def _initialize(self, use_gpu=False, enable_mkldnn=False):
  41. """
  42. initialize with the necessary elements
  43. """
  44. cfg = self.merge_configs()
  45. cfg.use_gpu = use_gpu
  46. if use_gpu:
  47. try:
  48. _places = os.environ["CUDA_VISIBLE_DEVICES"]
  49. int(_places[0])
  50. print("use gpu: ", use_gpu)
  51. print("CUDA_VISIBLE_DEVICES: ", _places)
  52. cfg.gpu_mem = 8000
  53. except:
  54. raise RuntimeError(
  55. "Environment Variable CUDA_VISIBLE_DEVICES is not set correctly. If you wanna use gpu, please set CUDA_VISIBLE_DEVICES via export CUDA_VISIBLE_DEVICES=cuda_device_id."
  56. )
  57. cfg.ir_optim = True
  58. cfg.enable_mkldnn = enable_mkldnn
  59. self.ser_predictor = SerPredictor(cfg)
  60. def merge_configs(self, ):
  61. # deafult cfg
  62. backup_argv = copy.deepcopy(sys.argv)
  63. sys.argv = sys.argv[:1]
  64. cfg = parse_args()
  65. update_cfg_map = vars(read_params())
  66. for key in update_cfg_map:
  67. cfg.__setattr__(key, update_cfg_map[key])
  68. sys.argv = copy.deepcopy(backup_argv)
  69. return cfg
  70. def read_images(self, paths=[]):
  71. images = []
  72. for img_path in paths:
  73. assert os.path.isfile(
  74. img_path), "The {} isn't a valid file.".format(img_path)
  75. img = cv2.imread(img_path)
  76. if img is None:
  77. logger.info("error in loading image:{}".format(img_path))
  78. continue
  79. images.append(img)
  80. return images
  81. def predict(self, images=[], paths=[]):
  82. """
  83. Get the chinese texts in the predicted images.
  84. Args:
  85. images (list(numpy.ndarray)): images data, shape of each is [H, W, C]. If images not paths
  86. paths (list[str]): The paths of images. If paths not images
  87. Returns:
  88. res (list): The result of chinese texts and save path of images.
  89. """
  90. if images != [] and isinstance(images, list) and paths == []:
  91. predicted_data = images
  92. elif images == [] and isinstance(paths, list) and paths != []:
  93. predicted_data = self.read_images(paths)
  94. else:
  95. raise TypeError("The input data is inconsistent with expectations.")
  96. assert predicted_data != [], "There is not any image to be predicted. Please check the input data."
  97. all_results = []
  98. for img in predicted_data:
  99. if img is None:
  100. logger.info("error in loading image")
  101. all_results.append([])
  102. continue
  103. starttime = time.time()
  104. ser_res, _, elapse = self.ser_predictor(img)
  105. elapse = time.time() - starttime
  106. logger.info("Predict time: {}".format(elapse))
  107. all_results.append(ser_res)
  108. return all_results
  109. @serving
  110. def serving_method(self, images, **kwargs):
  111. """
  112. Run as a service.
  113. """
  114. images_decode = [base64_to_cv2(image) for image in images]
  115. results = self.predict(images_decode, **kwargs)
  116. return results
  117. if __name__ == '__main__':
  118. ocr = OCRSystem()
  119. ocr._initialize()
  120. image_path = [
  121. './doc/imgs/11.jpg',
  122. './doc/imgs/12.jpg',
  123. ]
  124. res = ocr.predict(paths=image_path)
  125. print(res)