result.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403
  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 copy
  15. import math
  16. import os
  17. import subprocess
  18. import tempfile
  19. from pathlib import Path
  20. from typing import List, Optional
  21. import cv2
  22. import fitz
  23. import numpy as np
  24. import PIL
  25. from PIL import Image, ImageDraw, ImageFont
  26. from ....utils import logging
  27. from ....utils.file_interface import custom_open
  28. from ....utils.fonts import PINGFANG_FONT_FILE_PATH
  29. from ...common.result import BaseCVResult, JsonMixin
  30. class FormulaRecResult(BaseCVResult):
  31. def _get_input_fn(self):
  32. fn = super()._get_input_fn()
  33. if (page_idx := self["page_index"]) is not None:
  34. fp = Path(fn)
  35. stem, suffix = fp.stem, fp.suffix
  36. return f"{stem}_{page_idx}{suffix}"
  37. else:
  38. return fn
  39. def _to_str(self, *args, **kwargs):
  40. data = copy.deepcopy(self)
  41. data.pop("input_img")
  42. _str = JsonMixin._to_str(data, *args, **kwargs)["res"]
  43. return {"res": _str}
  44. def _to_json(self, *args, **kwargs):
  45. data = copy.deepcopy(self)
  46. data.pop("input_img")
  47. return JsonMixin._to_json(data, *args, **kwargs)
  48. def _to_img(
  49. self,
  50. ) -> Image.Image:
  51. """
  52. Draws a recognized formula on an image.
  53. This method processes an input image to recognize and render a LaTeX formula.
  54. It overlays the rendered formula onto the input image and returns the combined image.
  55. If the LaTeX rendering engine is not installed or a syntax error is detected,
  56. it logs a warning and returns the original image.
  57. Returns:
  58. Image.Image: An image with the recognized formula rendered alongside the original image.
  59. """
  60. image = Image.fromarray(self["input_img"])
  61. try:
  62. env_valid()
  63. except subprocess.CalledProcessError as e:
  64. logging.warning(
  65. "Please refer to 2.3 Formula Recognition Pipeline Visualization in Formula Recognition Pipeline Tutorial to install the LaTeX rendering engine at first."
  66. )
  67. return {"res": image}
  68. rec_formula = str(self["rec_formula"])
  69. image = np.array(image.convert("RGB"))
  70. xywh = crop_white_area(image)
  71. if xywh is not None:
  72. x, y, w, h = xywh
  73. image = image[y : y + h, x : x + w]
  74. image = Image.fromarray(image)
  75. image_width, image_height = image.size
  76. box = [[0, 0], [image_width, 0], [image_width, image_height], [0, image_height]]
  77. try:
  78. img_formula = draw_formula_module(
  79. image.size, box, rec_formula, is_debug=False
  80. )
  81. img_formula = Image.fromarray(img_formula)
  82. render_width, render_height = img_formula.size
  83. resize_height = render_height
  84. resize_width = int(resize_height * image_width / image_height)
  85. image = image.resize((resize_width, resize_height), Image.LANCZOS)
  86. new_image_width = image.width + int(render_width) + 10
  87. new_image = Image.new(
  88. "RGB", (new_image_width, render_height), (255, 255, 255)
  89. )
  90. new_image.paste(image, (0, 0))
  91. new_image.paste(img_formula, (image.width + 10, 0))
  92. return {"res": new_image}
  93. except subprocess.CalledProcessError as e:
  94. logging.warning("Syntax error detected in formula, rendering failed.")
  95. return {"res": image}
  96. def get_align_equation(equation: str) -> str:
  97. """
  98. Wraps an equation in LaTeX environment tags if not already aligned.
  99. This function checks if a given LaTeX equation contains any alignment tags (`align` or `align*`).
  100. If the equation does not contain these tags, it wraps the equation in `equation` and `nonumber` tags.
  101. Args:
  102. equation (str): The LaTeX equation to be checked and potentially modified.
  103. Returns:
  104. str: The modified equation with appropriate LaTeX tags for alignment.
  105. """
  106. is_align = False
  107. equation = str(equation) + "\n"
  108. begin_dict = [
  109. r"begin{align}",
  110. r"begin{align*}",
  111. ]
  112. for begin_sym in begin_dict:
  113. if begin_sym in equation:
  114. is_align = True
  115. break
  116. if not is_align:
  117. equation = (
  118. r"\begin{equation}"
  119. + "\n"
  120. + equation.strip()
  121. + r"\nonumber"
  122. + "\n"
  123. + r"\end{equation}"
  124. + "\n"
  125. )
  126. return equation
  127. def generate_tex_file(tex_file_path: str, equation: str) -> None:
  128. """
  129. Generates a LaTeX file containing a specific equation.
  130. This function creates a LaTeX file at the specified file path, writing the necessary
  131. LaTeX preamble and wrapping the provided equation in a document structure. The equation
  132. is processed to ensure it includes alignment tags if necessary.
  133. Args:
  134. tex_file_path (str): The file path where the LaTeX file will be saved.
  135. equation (str): The LaTeX equation to be written into the file.
  136. """
  137. with custom_open(tex_file_path, "w") as fp:
  138. start_template = (
  139. r"\documentclass{article}" + "\n"
  140. r"\usepackage{cite}" + "\n"
  141. r"\usepackage{amsmath,amssymb,amsfonts,upgreek}" + "\n"
  142. r"\usepackage{graphicx}" + "\n"
  143. r"\usepackage{textcomp}" + "\n"
  144. r"\DeclareMathSizes{14}{14}{9.8}{7}" + "\n"
  145. r"\pagestyle{empty}" + "\n"
  146. r"\begin{document}" + "\n"
  147. r"\begin{large}" + "\n"
  148. )
  149. fp.write(start_template)
  150. equation = get_align_equation(equation)
  151. fp.write(equation)
  152. end_template = r"\end{large}" + "\n" r"\end{document}" + "\n"
  153. fp.write(end_template)
  154. def generate_pdf_file(
  155. tex_path: str, pdf_dir: str, is_debug: bool = False
  156. ) -> Optional[bool]:
  157. """
  158. Generates a PDF file from a LaTeX file using pdflatex.
  159. This function checks if the specified LaTeX file exists, and then runs pdflatex to generate a PDF file
  160. in the specified directory. It can run in debug mode to show detailed output or in silent mode.
  161. Args:
  162. tex_path (str): The path to the LaTeX file.
  163. pdf_dir (str): The directory where the PDF file will be saved.
  164. is_debug (bool, optional): If True, runs pdflatex with detailed output. Defaults to False.
  165. Returns:
  166. Optional[bool]: Returns True if the PDF was generated successfully, False if the LaTeX file does not exist,
  167. and None if an error occurred during the pdflatex execution.
  168. """
  169. if os.path.exists(tex_path):
  170. command = "pdflatex -interaction=nonstopmode -halt-on-error -output-directory={} {}".format(
  171. pdf_dir, tex_path
  172. )
  173. if is_debug:
  174. subprocess.check_call(command, shell=True)
  175. else:
  176. custom_open(os.devnull, "w")
  177. subprocess.check_call(
  178. command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True
  179. )
  180. def crop_white_area(image: np.ndarray) -> Optional[List[int]]:
  181. """
  182. Finds and returns the bounding box of the non-white area in an image.
  183. This function converts an image to grayscale and uses binary thresholding to
  184. find contours. It then calculates the bounding rectangle around the non-white
  185. areas of the image.
  186. Args:
  187. image (np.ndarray): The input image as a NumPy array.
  188. Returns:
  189. Optional[List[int]]: A list [x, y, w, h] representing the bounding box of
  190. the non-white area, or None if no such area is found.
  191. """
  192. image = np.array(image).astype("uint8")
  193. gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  194. _, thresh = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)
  195. contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  196. if len(contours) > 0:
  197. x, y, w, h = cv2.boundingRect(np.concatenate(contours))
  198. return [x, y, w, h]
  199. else:
  200. return None
  201. def pdf2img(pdf_path: str, img_path: str, is_padding: bool = False):
  202. """
  203. Converts a single-page PDF to an image, optionally cropping white areas and adding padding.
  204. Args:
  205. pdf_path (str): The path to the PDF file.
  206. img_path (str): The path where the image will be saved.
  207. is_padding (bool): If True, adds a 30-pixel white padding around the image.
  208. Returns:
  209. np.ndarray: The resulting image as a NumPy array, or None if the PDF is not single-page.
  210. """
  211. pdfDoc = fitz.open(pdf_path)
  212. if pdfDoc.page_count != 1:
  213. return None
  214. for pg in range(pdfDoc.page_count):
  215. page = pdfDoc[pg]
  216. rotate = int(0)
  217. zoom_x = 2
  218. zoom_y = 2
  219. mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
  220. pix = page.get_pixmap(matrix=mat, alpha=False)
  221. getpngdata = pix.tobytes(output="png")
  222. # decode as np.uint8
  223. image_array = np.frombuffer(getpngdata, dtype=np.uint8)
  224. img = cv2.imdecode(image_array, cv2.IMREAD_ANYCOLOR)
  225. xywh = crop_white_area(img)
  226. if xywh is not None:
  227. x, y, w, h = xywh
  228. img = img[y : y + h, x : x + w]
  229. if is_padding:
  230. img = cv2.copyMakeBorder(
  231. img, 30, 30, 30, 30, cv2.BORDER_CONSTANT, value=(255, 255, 255)
  232. )
  233. return img
  234. return None
  235. def draw_formula_module(
  236. img_size: tuple, box: list, formula: str, is_debug: bool = False
  237. ):
  238. """
  239. Draw box formula for module.
  240. Args:
  241. img_size (tuple): The size of the image as (width, height).
  242. box (list): The coordinates for the bounding box.
  243. formula (str): The LaTeX formula to render.
  244. is_debug (bool): If True, retains intermediate files for debugging purposes.
  245. Returns:
  246. np.ndarray: The resulting image with the formula or an error message.
  247. """
  248. box_width, box_height = img_size
  249. with tempfile.TemporaryDirectory() as td:
  250. tex_file_path = os.path.join(td, "temp.tex")
  251. pdf_file_path = os.path.join(td, "temp.pdf")
  252. img_file_path = os.path.join(td, "temp.jpg")
  253. generate_tex_file(tex_file_path, formula)
  254. if os.path.exists(tex_file_path):
  255. generate_pdf_file(tex_file_path, td, is_debug)
  256. formula_img = None
  257. if os.path.exists(pdf_file_path):
  258. formula_img = pdf2img(pdf_file_path, img_file_path, is_padding=False)
  259. if formula_img is not None:
  260. return formula_img
  261. else:
  262. img_right_text = draw_box_txt_fine(
  263. img_size, box, "Rendering Failed", PINGFANG_FONT_FILE_PATH
  264. )
  265. return img_right_text
  266. def env_valid() -> bool:
  267. """
  268. Validates if the environment is correctly set up to convert LaTeX formulas to images.
  269. Returns:
  270. bool: True if the environment is valid and the conversion is successful, False otherwise.
  271. """
  272. with tempfile.TemporaryDirectory() as td:
  273. tex_file_path = os.path.join(td, "temp.tex")
  274. pdf_file_path = os.path.join(td, "temp.pdf")
  275. img_file_path = os.path.join(td, "temp.jpg")
  276. formula = "a+b=c"
  277. is_debug = False
  278. generate_tex_file(tex_file_path, formula)
  279. if os.path.exists(tex_file_path):
  280. generate_pdf_file(tex_file_path, td, is_debug)
  281. if os.path.exists(pdf_file_path):
  282. formula_img = pdf2img(pdf_file_path, img_file_path, is_padding=False)
  283. def draw_box_txt_fine(img_size: tuple, box: list, txt: str, font_path: str):
  284. """
  285. Draw box text.
  286. Args:
  287. img_size (tuple): Size of the image as (width, height).
  288. box (list): List of four points defining the box, each point is a tuple (x, y).
  289. txt (str): The text to draw inside the box.
  290. font_path (str): Path to the font file to be used for drawing text.
  291. Returns:
  292. np.ndarray: Image array with the text drawn and transformed to fit the box.
  293. """
  294. box_height = int(
  295. math.sqrt((box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2)
  296. )
  297. box_width = int(
  298. math.sqrt((box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2)
  299. )
  300. if box_height > 2 * box_width and box_height > 30:
  301. img_text = Image.new("RGB", (box_height, box_width), (255, 255, 255))
  302. draw_text = ImageDraw.Draw(img_text)
  303. if txt:
  304. font = create_font(txt, (box_height, box_width), font_path)
  305. draw_text.text([0, 0], txt, fill=(0, 0, 0), font=font)
  306. img_text = img_text.transpose(Image.ROTATE_270)
  307. else:
  308. img_text = Image.new("RGB", (box_width, box_height), (255, 255, 255))
  309. draw_text = ImageDraw.Draw(img_text)
  310. if txt:
  311. font = create_font(txt, (box_width, box_height), font_path)
  312. draw_text.text([0, 0], txt, fill=(0, 0, 0), font=font)
  313. pts1 = np.float32(
  314. [[0, 0], [box_width, 0], [box_width, box_height], [0, box_height]]
  315. )
  316. pts2 = np.array(box, dtype=np.float32)
  317. M = cv2.getPerspectiveTransform(pts1, pts2)
  318. img_text = np.array(img_text, dtype=np.uint8)
  319. img_right_text = cv2.warpPerspective(
  320. img_text,
  321. M,
  322. img_size,
  323. flags=cv2.INTER_NEAREST,
  324. borderMode=cv2.BORDER_CONSTANT,
  325. borderValue=(255, 255, 255),
  326. )
  327. return img_right_text
  328. def create_font(txt: str, sz: tuple, font_path: str) -> ImageFont.FreeTypeFont:
  329. """
  330. Creates a font object with a size that ensures the text fits within the specified dimensions.
  331. Args:
  332. txt (str): The text to fit.
  333. sz (tuple): The target size as (width, height).
  334. font_path (str): The path to the font file.
  335. Returns:
  336. ImageFont.FreeTypeFont: A PIL font object at the appropriate size.
  337. """
  338. font_size = int(sz[1] * 0.8)
  339. font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
  340. if int(PIL.__version__.split(".")[0]) < 10:
  341. length = font.getsize(txt)[0]
  342. else:
  343. length = font.getlength(txt)
  344. if length > sz[0]:
  345. font_size = int(font_size * sz[0] / length)
  346. font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
  347. return font