module.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143
  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.table.predict_table import TableSystem as _TableSystem
  30. from ppstructure.predict_system import save_structure_res
  31. from ppstructure.utility import parse_args
  32. from deploy.hubserving.structure_table.params import read_params
  33. @moduleinfo(
  34. name="structure_table",
  35. version="1.0.0",
  36. summary="PP-Structure table service",
  37. author="paddle-dev",
  38. author_email="paddle-dev@baidu.com",
  39. type="cv/structure_table")
  40. class TableSystem(hub.Module):
  41. def _initialize(self, use_gpu=False, enable_mkldnn=False):
  42. """
  43. initialize with the necessary elements
  44. """
  45. cfg = self.merge_configs()
  46. cfg.use_gpu = use_gpu
  47. if use_gpu:
  48. try:
  49. _places = os.environ["CUDA_VISIBLE_DEVICES"]
  50. int(_places[0])
  51. print("use gpu: ", use_gpu)
  52. print("CUDA_VISIBLE_DEVICES: ", _places)
  53. cfg.gpu_mem = 8000
  54. except:
  55. raise RuntimeError(
  56. "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."
  57. )
  58. cfg.ir_optim = True
  59. cfg.enable_mkldnn = enable_mkldnn
  60. self.table_sys = _TableSystem(cfg)
  61. def merge_configs(self):
  62. # deafult cfg
  63. backup_argv = copy.deepcopy(sys.argv)
  64. sys.argv = sys.argv[:1]
  65. cfg = parse_args()
  66. update_cfg_map = vars(read_params())
  67. for key in update_cfg_map:
  68. cfg.__setattr__(key, update_cfg_map[key])
  69. sys.argv = copy.deepcopy(backup_argv)
  70. return cfg
  71. def read_images(self, paths=[]):
  72. images = []
  73. for img_path in paths:
  74. assert os.path.isfile(
  75. img_path), "The {} isn't a valid file.".format(img_path)
  76. img = cv2.imread(img_path)
  77. if img is None:
  78. logger.info("error in loading image:{}".format(img_path))
  79. continue
  80. images.append(img)
  81. return images
  82. def predict(self, images=[], paths=[]):
  83. """
  84. Get the chinese texts in the predicted images.
  85. Args:
  86. images (list(numpy.ndarray)): images data, shape of each is [H, W, C]. If images not paths
  87. paths (list[str]): The paths of images. If paths not images
  88. Returns:
  89. res (list): The result of chinese texts and save path of images.
  90. """
  91. if images != [] and isinstance(images, list) and paths == []:
  92. predicted_data = images
  93. elif images == [] and isinstance(paths, list) and paths != []:
  94. predicted_data = self.read_images(paths)
  95. else:
  96. raise TypeError("The input data is inconsistent with expectations.")
  97. assert predicted_data != [], "There is not any image to be predicted. Please check the input data."
  98. all_results = []
  99. for img in predicted_data:
  100. if img is None:
  101. logger.info("error in loading image")
  102. all_results.append([])
  103. continue
  104. starttime = time.time()
  105. res, _ = self.table_sys(img)
  106. elapse = time.time() - starttime
  107. logger.info("Predict time: {}".format(elapse))
  108. all_results.append({'html': res['html']})
  109. return all_results
  110. @serving
  111. def serving_method(self, images, **kwargs):
  112. """
  113. Run as a service.
  114. """
  115. images_decode = [base64_to_cv2(image) for image in images]
  116. results = self.predict(images_decode, **kwargs)
  117. return results
  118. if __name__ == '__main__':
  119. table_system = TableSystem()
  120. table_system._initialize()
  121. image_path = ['./ppstructure/docs/table/table.jpg']
  122. res = table_system.predict(paths=image_path)
  123. print(res)