table_recognition.py 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  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 tempfile
  15. from typing import Any, Dict, List
  16. import ultra_infer as ui
  17. import numpy as np
  18. from paddlex.inference.common.batch_sampler import ImageBatchSampler
  19. from paddlex.inference.models.table_structure_recognition.result import (
  20. TableRecResult,
  21. )
  22. from paddlex.modules.table_recognition.model_list import MODELS
  23. from paddlex_hpi._utils.compat import get_compat_version
  24. from paddlex_hpi.models.base import CVPredictor
  25. class TablePredictor(CVPredictor):
  26. entities = MODELS
  27. def _build_ui_model(
  28. self, option: ui.RuntimeOption
  29. ) -> ui.vision.ocr.StructureV2Table:
  30. compat_version = get_compat_version()
  31. if compat_version == "2.5" or self.model_name == "SLANet":
  32. bbox_shape_type = "ori"
  33. else:
  34. bbox_shape_type = "pad"
  35. with tempfile.NamedTemporaryFile("w", encoding="utf-8", suffix=".txt") as f:
  36. pp_config = self.config["PostProcess"]
  37. for lab in pp_config["character_dict"]:
  38. f.write(lab + "\n")
  39. f.flush()
  40. model = ui.vision.ocr.StructureV2Table(
  41. str(self.model_path),
  42. str(self.params_path),
  43. table_char_dict_path=f.name,
  44. box_shape=bbox_shape_type,
  45. runtime_option=option,
  46. )
  47. return model
  48. def _build_batch_sampler(self) -> ImageBatchSampler:
  49. return ImageBatchSampler()
  50. def _get_result_class(self) -> type:
  51. return TableRecResult
  52. def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
  53. batch_raw_imgs = self._data_reader(imgs=batch_data.instances)
  54. imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
  55. ui_results = self._ui_model.batch_predict(imgs)
  56. bbox_list = []
  57. structure_list = []
  58. structure_score_list = []
  59. for ui_result in ui_results:
  60. bbox_list.append(ui_result.table_boxes)
  61. structure_list.append(ui_result.table_structure)
  62. structure_score_list.append(0.0)
  63. return {
  64. "input_path": batch_data.input_paths,
  65. "page_index": batch_data.page_indexes,
  66. "input_img": batch_raw_imgs,
  67. "bbox": bbox_list,
  68. "structure": structure_list,
  69. "structure_score": structure_score_list,
  70. }