result.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287
  1. # Copyright (c) 2024 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 math
  15. import os
  16. import random
  17. import subprocess
  18. import tempfile
  19. from pathlib import Path
  20. from typing import Dict, Tuple
  21. import cv2
  22. import numpy as np
  23. from PIL import Image, ImageDraw
  24. from ....utils import logging
  25. from ....utils.fonts import PINGFANG_FONT_FILE_PATH
  26. from ...common.result import BaseCVResult, JsonMixin
  27. from ...models.formula_recognition.result import (
  28. crop_white_area,
  29. draw_box_txt_fine,
  30. draw_formula_module,
  31. env_valid,
  32. generate_pdf_file,
  33. generate_tex_file,
  34. pdf2img,
  35. )
  36. class FormulaRecognitionResult(BaseCVResult):
  37. """Formula Recognition Result"""
  38. def _get_input_fn(self):
  39. fn = super()._get_input_fn()
  40. if (page_idx := self["page_index"]) is not None:
  41. fp = Path(fn)
  42. stem, suffix = fp.stem, fp.suffix
  43. return f"{stem}_{page_idx}{suffix}"
  44. else:
  45. return fn
  46. def _to_img(self) -> Dict[str, Image.Image]:
  47. """
  48. Converts the internal data to a PIL Image with detection and recognition results.
  49. Returns:
  50. Dict[str, Image.Image]: An image with detection boxes, texts, and scores blended on it.
  51. """
  52. image = Image.fromarray(self["doc_preprocessor_res"]["output_img"])
  53. res_img_dict = {}
  54. model_settings = self["model_settings"]
  55. if model_settings["use_doc_preprocessor"]:
  56. res_img_dict.update(**self["doc_preprocessor_res"].img)
  57. layout_det_res = self["layout_det_res"]
  58. if len(layout_det_res) > 0:
  59. res_img_dict["layout_det_res"] = layout_det_res.img["res"]
  60. try:
  61. env_valid()
  62. except subprocess.CalledProcessError as e:
  63. logging.warning(
  64. "Please refer to 2.3 Formula Recognition Pipeline Visualization in Formula Recognition Pipeline Tutorial to install the LaTeX rendering engine at first."
  65. )
  66. res_img_dict["formula_res_img"] = image
  67. return res_img_dict
  68. if len(layout_det_res) <= 0:
  69. image = np.array(image.convert("RGB"))
  70. rec_formula = self["formula_res_list"][0]["rec_formula"]
  71. xywh = crop_white_area(image)
  72. if xywh is not None:
  73. x, y, w, h = xywh
  74. image = image[y : y + h, x : x + w]
  75. image = Image.fromarray(image)
  76. image_width, image_height = image.size
  77. box = [
  78. [0, 0],
  79. [image_width, 0],
  80. [image_width, image_height],
  81. [0, image_height],
  82. ]
  83. try:
  84. img_formula = draw_formula_module(
  85. image.size, box, rec_formula, is_debug=False
  86. )
  87. img_formula = Image.fromarray(img_formula)
  88. render_width, render_height = img_formula.size
  89. resize_height = render_height
  90. resize_width = int(resize_height * image_width / image_height)
  91. image = image.resize((resize_width, resize_height), Image.LANCZOS)
  92. new_image_width = image.width + int(render_width) + 10
  93. new_image = Image.new(
  94. "RGB", (new_image_width, render_height), (255, 255, 255)
  95. )
  96. new_image.paste(image, (0, 0))
  97. new_image.paste(img_formula, (image.width + 10, 0))
  98. res_img_dict["formula_res_img"] = new_image
  99. return res_img_dict
  100. except subprocess.CalledProcessError as e:
  101. logging.warning("Syntax error detected in formula, rendering failed.")
  102. res_img_dict["formula_res_img"] = image
  103. return res_img_dict
  104. h, w = image.height, image.width
  105. img_left = image.copy()
  106. img_right = np.ones((h, w, 3), dtype=np.uint8) * 255
  107. random.seed(0)
  108. draw_left = ImageDraw.Draw(img_left)
  109. self["formula_res_list"]
  110. for tno in range(len(self["formula_res_list"])):
  111. formula_res = self["formula_res_list"][tno]
  112. formula_res["formula_region_id"]
  113. formula = str(formula_res["rec_formula"])
  114. dt_polys = formula_res["dt_polys"]
  115. x1, y1, x2, y2 = list(dt_polys)
  116. try:
  117. color = (
  118. random.randint(0, 255),
  119. random.randint(0, 255),
  120. random.randint(0, 255),
  121. )
  122. box = [x1, y1, x2, y1, x2, y2, x1, y2]
  123. box = np.array(box).reshape([-1, 2])
  124. pts = [(x, y) for x, y in box.tolist()]
  125. draw_left.polygon(pts, outline=color, width=8)
  126. draw_left.polygon(box, fill=color)
  127. img_right_text = draw_box_formula_fine(
  128. (w, h),
  129. box,
  130. formula,
  131. is_debug=False,
  132. )
  133. pts = np.array(box, np.int32).reshape((-1, 1, 2))
  134. cv2.polylines(img_right_text, [pts], True, color, 1)
  135. img_right = cv2.bitwise_and(img_right, img_right_text)
  136. except subprocess.CalledProcessError as e:
  137. logging.warning("Syntax error detected in formula, rendering failed.")
  138. continue
  139. img_left = Image.blend(image, img_left, 0.5)
  140. img_show = Image.new("RGB", (int(w * 2), h), (255, 255, 255))
  141. img_show.paste(img_left, (0, 0, w, h))
  142. img_show.paste(Image.fromarray(img_right), (w, 0, w * 2, h))
  143. res_img_dict["formula_res_img"] = img_show
  144. return res_img_dict
  145. def _to_str(self, *args, **kwargs) -> Dict[str, str]:
  146. """Converts the instance's attributes to a dictionary and then to a string.
  147. Args:
  148. *args: Additional positional arguments passed to the base class method.
  149. **kwargs: Additional keyword arguments passed to the base class method.
  150. Returns:
  151. Dict[str, str]: A dictionary with the instance's attributes converted to strings.
  152. """
  153. data = {}
  154. data["input_path"] = self["input_path"]
  155. data["page_index"] = self["page_index"]
  156. data["model_settings"] = self["model_settings"]
  157. if self["model_settings"]["use_doc_preprocessor"]:
  158. data["doc_preprocessor_res"] = self["doc_preprocessor_res"].str["res"]
  159. if len(self["layout_det_res"]) > 0:
  160. data["layout_det_res"] = self["layout_det_res"].str["res"]
  161. data["formula_res_list"] = []
  162. for tno in range(len(self["formula_res_list"])):
  163. rec_formula_dict = {
  164. "rec_formula": self["formula_res_list"][tno]["rec_formula"],
  165. "formula_region_id": self["formula_res_list"][tno]["formula_region_id"],
  166. }
  167. if "dt_polys" in self["formula_res_list"][tno]:
  168. rec_formula_dict["dt_polys"] = (
  169. self["formula_res_list"][tno]["dt_polys"],
  170. )
  171. data["formula_res_list"].append(rec_formula_dict)
  172. return JsonMixin._to_str(data, *args, **kwargs)
  173. def _to_json(self, *args, **kwargs) -> Dict[str, str]:
  174. """
  175. Converts the object's data to a JSON dictionary.
  176. Args:
  177. *args: Positional arguments passed to the JsonMixin._to_json method.
  178. **kwargs: Keyword arguments passed to the JsonMixin._to_json method.
  179. Returns:
  180. Dict[str, str]: A dictionary containing the object's data in JSON format.
  181. """
  182. data = {}
  183. data["input_path"] = self["input_path"]
  184. data["page_index"] = str(self["page_index"])
  185. data["model_settings"] = self["model_settings"]
  186. if self["model_settings"]["use_doc_preprocessor"]:
  187. data["doc_preprocessor_res"] = self["doc_preprocessor_res"].str["res"]
  188. if len(self["layout_det_res"]) > 0:
  189. data["layout_det_res"] = self["layout_det_res"].str["res"]
  190. data["formula_res_list"] = []
  191. for tno in range(len(self["formula_res_list"])):
  192. rec_formula_dict = {
  193. "rec_formula": self["formula_res_list"][tno]["rec_formula"],
  194. "formula_region_id": self["formula_res_list"][tno]["formula_region_id"],
  195. }
  196. if "dt_polys" in self["formula_res_list"][tno]:
  197. rec_formula_dict["dt_polys"] = (
  198. self["formula_res_list"][tno]["dt_polys"],
  199. )
  200. data["formula_res_list"].append(rec_formula_dict)
  201. return JsonMixin._to_json(data, *args, **kwargs)
  202. def draw_box_formula_fine(
  203. img_size: Tuple[int, int], box: np.ndarray, formula: str, is_debug: bool = False
  204. ) -> np.ndarray:
  205. """draw box formula for pipeline"""
  206. """
  207. Draw box formula for pipeline.
  208. This function generates a LaTeX formula image and transforms it to fit
  209. within a specified bounding box on a larger image. If the rendering fails,
  210. it will write "Rendering Failed" inside the box.
  211. Args:
  212. img_size (Tuple[int, int]): The size of the image (width, height).
  213. box (np.ndarray): A numpy array representing the four corners of the bounding box.
  214. formula (str): The LaTeX formula to render.
  215. is_debug (bool, optional): If True, enables debug mode. Defaults to False.
  216. Returns:
  217. np.ndarray: An image array with the rendered formula inside the specified box.
  218. """
  219. box_height = int(
  220. math.sqrt((box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2)
  221. )
  222. box_width = int(
  223. math.sqrt((box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2)
  224. )
  225. with tempfile.TemporaryDirectory() as td:
  226. tex_file_path = os.path.join(td, "temp.tex")
  227. pdf_file_path = os.path.join(td, "temp.pdf")
  228. img_file_path = os.path.join(td, "temp.jpg")
  229. generate_tex_file(tex_file_path, formula)
  230. if os.path.exists(tex_file_path):
  231. generate_pdf_file(tex_file_path, td, is_debug)
  232. formula_img = None
  233. if os.path.exists(pdf_file_path):
  234. formula_img = pdf2img(pdf_file_path, img_file_path, is_padding=False)
  235. if formula_img is not None:
  236. formula_h, formula_w = formula_img.shape[:-1]
  237. resize_height = box_height
  238. resize_width = formula_w * resize_height / formula_h
  239. formula_img = cv2.resize(
  240. formula_img, (int(resize_width), int(resize_height))
  241. )
  242. formula_h, formula_w = formula_img.shape[:-1]
  243. pts1 = np.float32(
  244. [[0, 0], [box_width, 0], [box_width, box_height], [0, box_height]]
  245. )
  246. pts2 = np.array(box, dtype=np.float32)
  247. M = cv2.getPerspectiveTransform(pts1, pts2)
  248. formula_img = np.array(formula_img, dtype=np.uint8)
  249. img_right_text = cv2.warpPerspective(
  250. formula_img,
  251. M,
  252. img_size,
  253. flags=cv2.INTER_NEAREST,
  254. borderMode=cv2.BORDER_CONSTANT,
  255. borderValue=(255, 255, 255),
  256. )
  257. else:
  258. img_right_text = draw_box_txt_fine(
  259. img_size, box, "Rendering Failed", PINGFANG_FONT_FILE_PATH
  260. )
  261. return img_right_text