table_structure.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
  1. # Copyright (c) 2020 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. import time
  15. from typing import Any, Dict, List, Tuple
  16. import numpy as np
  17. from mineru.utils.os_env_config import get_op_num_threads
  18. from .table_structure_utils import (
  19. OrtInferSession,
  20. TableLabelDecode,
  21. TablePreprocess,
  22. BatchTablePreprocess,
  23. )
  24. class TableStructurer:
  25. def __init__(self, config: Dict[str, Any]):
  26. self.preprocess_op = TablePreprocess()
  27. self.batch_preprocess_op = BatchTablePreprocess()
  28. config["intra_op_num_threads"] = get_op_num_threads("MINERU_INTRA_OP_NUM_THREADS")
  29. config["inter_op_num_threads"] = get_op_num_threads("MINERU_INTER_OP_NUM_THREADS")
  30. self.session = OrtInferSession(config)
  31. self.character = self.session.get_metadata()
  32. self.postprocess_op = TableLabelDecode(self.character)
  33. def process(self, img):
  34. starttime = time.time()
  35. data = {"image": img}
  36. data = self.preprocess_op(data)
  37. img = data[0]
  38. if img is None:
  39. return None, 0
  40. img = np.expand_dims(img, axis=0)
  41. img = img.copy()
  42. outputs = self.session([img])
  43. preds = {"loc_preds": outputs[0], "structure_probs": outputs[1]}
  44. shape_list = np.expand_dims(data[-1], axis=0)
  45. post_result = self.postprocess_op(preds, [shape_list])
  46. bbox_list = post_result["bbox_batch_list"][0]
  47. structure_str_list = post_result["structure_batch_list"][0]
  48. structure_str_list = structure_str_list[0]
  49. structure_str_list = (
  50. ["<html>", "<body>", "<table>"]
  51. + structure_str_list
  52. + ["</table>", "</body>", "</html>"]
  53. )
  54. elapse = time.time() - starttime
  55. return structure_str_list, bbox_list, elapse
  56. def batch_process(
  57. self, img_list: List[np.ndarray]
  58. ) -> List[Tuple[List[str], np.ndarray, float]]:
  59. """批量处理图像列表
  60. Args:
  61. img_list: 图像列表
  62. Returns:
  63. 结果列表,每个元素包含 (table_struct_str, cell_bboxes, elapse)
  64. """
  65. starttime = time.perf_counter()
  66. batch_data = self.batch_preprocess_op(img_list)
  67. preprocessed_images = batch_data[0]
  68. shape_lists = batch_data[1]
  69. preprocessed_images = np.array(preprocessed_images)
  70. bbox_preds, struct_probs = self.session([preprocessed_images])
  71. batch_size = preprocessed_images.shape[0]
  72. results = []
  73. for bbox_pred, struct_prob, shape_list in zip(
  74. bbox_preds, struct_probs, shape_lists
  75. ):
  76. preds = {
  77. "loc_preds": np.expand_dims(bbox_pred, axis=0),
  78. "structure_probs": np.expand_dims(struct_prob, axis=0),
  79. }
  80. shape_list = np.expand_dims(shape_list, axis=0)
  81. post_result = self.postprocess_op(preds, [shape_list])
  82. bbox_list = post_result["bbox_batch_list"][0]
  83. structure_str_list = post_result["structure_batch_list"][0]
  84. structure_str_list = structure_str_list[0]
  85. structure_str_list = (
  86. ["<html>", "<body>", "<table>"]
  87. + structure_str_list
  88. + ["</table>", "</body>", "</html>"]
  89. )
  90. results.append((structure_str_list, bbox_list, 0))
  91. total_elapse = time.perf_counter() - starttime
  92. for i in range(len(results)):
  93. results[i] = (results[i][0], results[i][1], total_elapse / batch_size)
  94. return results