table_rec.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  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 cv2
  15. import numpy as np
  16. from pathlib import Path
  17. from .utils.mixin import HtmlMixin, XlsxMixin
  18. from .base import BaseResult, CVResult
  19. class TableRecResult(CVResult, HtmlMixin):
  20. """SaveTableResults"""
  21. def __init__(self, data):
  22. super().__init__(data)
  23. HtmlMixin.__init__(self)
  24. self._show_func_register("save_to_html")(self.save_to_html)
  25. def _to_html(self):
  26. return self["html"]
  27. def _to_img(self):
  28. image = self._img_reader.read(self["input_path"])
  29. bbox_res = self["bbox"]
  30. if len(bbox_res) > 0 and len(bbox_res[0]) == 4:
  31. vis_img = self.draw_rectangle(image, bbox_res)
  32. else:
  33. vis_img = self.draw_bbox(image, bbox_res)
  34. return vis_img
  35. def draw_rectangle(self, image, boxes):
  36. """draw_rectangle"""
  37. boxes = np.array(boxes)
  38. img_show = image.copy()
  39. for box in boxes.astype(int):
  40. x1, y1, x2, y2 = box
  41. cv2.rectangle(img_show, (x1, y1), (x2, y2), (255, 0, 0), 2)
  42. return img_show
  43. def draw_bbox(self, image, boxes):
  44. """draw_bbox"""
  45. for box in boxes:
  46. box = np.reshape(np.array(box), [-1, 1, 2]).astype(np.int64)
  47. image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2)
  48. return image
  49. class StructureTableResult(TableRecResult, XlsxMixin):
  50. """StructureTableResult"""
  51. def __init__(self, data):
  52. super().__init__(data)
  53. XlsxMixin.__init__(self)
  54. class TableResult(BaseResult):
  55. """TableResult"""
  56. def save_to_html(self, save_path):
  57. if not save_path.lower().endswith(("html")):
  58. input_path = self["input_path"]
  59. save_path = Path(save_path) / f"{Path(input_path).stem}"
  60. else:
  61. save_path = Path(save_path).stem
  62. for table_result in self["table_result"]:
  63. table_result.save_to_html(save_path)
  64. def save_to_xlsx(self, save_path):
  65. if not save_path.lower().endswith(("xlsx")):
  66. input_path = self["input_path"]
  67. save_path = Path(save_path) / f"{Path(input_path).stem}"
  68. else:
  69. save_path = Path(save_path).stem
  70. for table_result in self["table_result"]:
  71. table_result.save_to_xlsx(save_path)
  72. def save_to_img(self, save_path):
  73. if not save_path.lower().endswith((".jpg", ".png")):
  74. input_path = self["input_path"]
  75. save_path = Path(save_path) / f"{Path(input_path).stem}"
  76. else:
  77. save_path = Path(save_path).stem
  78. layout_save_path = f"{save_path}_layout.jpg"
  79. ocr_save_path = f"{save_path}_ocr.jpg"
  80. table_save_path = f"{save_path}_table.jpg"
  81. layout_result = self["layout_result"]
  82. layout_result.save_to_img(layout_save_path)
  83. ocr_result = self["ocr_result"]
  84. ocr_result.save_to_img(ocr_save_path)
  85. for table_result in self["table_result"]:
  86. table_result.save_to_img(table_save_path)