chat_ocr.py 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159
  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 copy
  15. from pathlib import Path
  16. from .base import BaseResult
  17. from .utils.mixin import Base64Mixin
  18. class LayoutParsingResult(BaseResult):
  19. """LayoutParsingResult"""
  20. pass
  21. class VisualInfoResult(BaseResult):
  22. """VisualInfoResult"""
  23. pass
  24. class VisualResult(BaseResult):
  25. """VisualInfoResult"""
  26. def __init__(self, data, page_id=None, src_input_name=None):
  27. super().__init__(data)
  28. self.page_id = page_id
  29. self.src_input_name = src_input_name
  30. def _to_str(self, _, *args, **kwargs):
  31. return super()._to_str(
  32. {"layout_parsing_result": self["layout_parsing_result"]}, *args, **kwargs
  33. )
  34. def get_target_name(self, save_path):
  35. if self.src_input_name.endswith(".pdf"):
  36. save_path = (
  37. Path(save_path)
  38. / f"{Path(self.src_input_name).stem}_pdf"
  39. / Path("page_{:04d}".format(self.page_id + 1))
  40. )
  41. else:
  42. save_path = Path(save_path) / f"{Path(self.src_input_name).stem}"
  43. return save_path
  44. def save_to_json(self, save_path):
  45. if not save_path.lower().endswith(("json")):
  46. save_path = self.get_target_name(save_path)
  47. else:
  48. save_path = Path(save_path).stem
  49. layout_save_path = f"{save_path}_layout.jpg"
  50. ocr_save_path = f"{save_path}_ocr.jpg"
  51. table_save_path = f"{save_path}_table"
  52. self["input_path"] = layout_save_path
  53. self["layout_result"]["input_path"] = layout_save_path
  54. self["layout_parsing_result"]["input_path"] = layout_save_path
  55. self["ocr_result"]["input_path"] = ocr_save_path
  56. table_result_num = len(self["table_result"])
  57. for idx in range(table_result_num):
  58. self["table_result"][idx]["input_path"] = "{}_{:04d}".format(
  59. table_save_path, idx + 1
  60. )
  61. if not str(save_path).endswith(".json"):
  62. save_path = "{}.json".format(save_path)
  63. super().save_to_json(save_path)
  64. def save_to_html(self, save_path):
  65. if not save_path.lower().endswith(("html")):
  66. save_path = self.get_target_name(save_path)
  67. else:
  68. save_path = Path(save_path).stem
  69. table_save_path = f"{save_path}_table"
  70. for idx, table_result in enumerate(self["table_result"]):
  71. basename = (Path(table_result["input_path"]).name).split(".")[0]
  72. table_result["input_path"] = Path(
  73. str(table_result["input_path"]).replace(
  74. basename, "{}_{:04d}".format(table_save_path, idx + 1)
  75. )
  76. )
  77. table_result.save_to_html(save_path)
  78. def save_to_xlsx(self, save_path):
  79. if not save_path.lower().endswith(("xlsx")):
  80. save_path = self.get_target_name(save_path)
  81. else:
  82. save_path = Path(save_path).stem
  83. table_save_path = f"{save_path}_table"
  84. for idx, table_result in enumerate(self["table_result"]):
  85. basename = (Path(table_result["input_path"]).name).split(".")[0]
  86. table_result["input_path"] = Path(
  87. str(table_result["input_path"]).replace(
  88. basename, "{}_{:04d}".format(table_save_path, idx + 1)
  89. )
  90. )
  91. table_result.save_to_xlsx(save_path)
  92. def save_to_img(self, save_path):
  93. if not save_path.lower().endswith((".jpg", ".png")):
  94. save_path = self.get_target_name(save_path)
  95. else:
  96. save_path = Path(save_path).stem
  97. oricls_save_path = f"{save_path}_oricls.jpg"
  98. oricls_result = self["oricls_result"]
  99. if oricls_result:
  100. oricls_result.save_to_img(oricls_save_path)
  101. uvdoc_save_path = f"{save_path}_uvdoc.jpg"
  102. unwarp_result = self["unwarp_result"]
  103. if unwarp_result:
  104. unwarp_result.save_to_img(uvdoc_save_path)
  105. curve_save_path = f"{save_path}_curve"
  106. curve_results = self["curve_result"]
  107. # TODO(): support list of result
  108. if isinstance(curve_results, dict):
  109. curve_results = [curve_results]
  110. for idx, curve_result in enumerate(curve_results):
  111. curve_result.save_to_img(f"{curve_save_path}_{idx}.jpg")
  112. layout_save_path = f"{save_path}_layout.jpg"
  113. layout_result = self["layout_result"]
  114. if layout_result:
  115. layout_result.save_to_img(layout_save_path)
  116. ocr_save_path = f"{save_path}_ocr.jpg"
  117. table_save_path = f"{save_path}_table"
  118. ocr_result = self["ocr_result"]
  119. if ocr_result:
  120. ocr_result.save_to_img(ocr_save_path)
  121. for idx, table_result in enumerate(self["table_result"]):
  122. table_result.save_to_img(f"{table_save_path}_{idx}.jpg")
  123. class VectorResult(BaseResult, Base64Mixin):
  124. """VisualInfoResult"""
  125. def _to_base64(self):
  126. return self["vector"]
  127. class RetrievalResult(BaseResult):
  128. """VisualInfoResult"""
  129. pass
  130. class ChatResult(BaseResult):
  131. """VisualInfoResult"""
  132. pass