formula_rec.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365
  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 os
  15. import sys
  16. import cv2
  17. import math
  18. import random
  19. import tempfile
  20. import subprocess
  21. import numpy as np
  22. from pathlib import Path
  23. from PIL import Image, ImageDraw
  24. from .base import BaseResult, CVResult
  25. from ...utils import logging
  26. from .ocr import draw_box_txt_fine
  27. from ...utils.fonts import PINGFANG_FONT_FILE_PATH
  28. class FormulaRecResult(CVResult):
  29. def _to_str(self, *args, **kwargs):
  30. return super()._to_str(*args, **kwargs).replace("\\\\", "\\")
  31. def _to_img(
  32. self,
  33. ):
  34. """Draw formula on image"""
  35. try:
  36. env_valid()
  37. except subprocess.CalledProcessError as e:
  38. logging.warning(
  39. "Please refer to 2.3 Formula Recognition Pipeline Visualization in Formula Recognition Pipeline Tutorial to install the LaTeX rendering engine at first."
  40. )
  41. return None
  42. image = self._img_reader.read(self["input_path"])
  43. rec_formula = str(self["rec_text"])
  44. image = np.array(image.convert("RGB"))
  45. xywh = crop_white_area(image)
  46. if xywh is not None:
  47. x, y, w, h = xywh
  48. image = image[y : y + h, x : x + w]
  49. image = Image.fromarray(image)
  50. image_width, image_height = image.size
  51. box = [[0, 0], [image_width, 0], [image_width, image_height], [0, image_height]]
  52. try:
  53. img_formula = draw_formula_module(
  54. image.size, box, rec_formula, is_debug=False
  55. )
  56. img_formula = Image.fromarray(img_formula)
  57. render_width, render_height = img_formula.size
  58. resize_height = render_height
  59. resize_width = int(resize_height * image_width / image_height)
  60. image = image.resize((resize_width, resize_height), Image.LANCZOS)
  61. new_image_width = image.width + int(render_width) + 10
  62. new_image = Image.new(
  63. "RGB", (new_image_width, render_height), (255, 255, 255)
  64. )
  65. new_image.paste(image, (0, 0))
  66. new_image.paste(img_formula, (image.width + 10, 0))
  67. return new_image
  68. except subprocess.CalledProcessError as e:
  69. logging.warning("Syntax error detected in formula, rendering failed.")
  70. return None
  71. class FormulaResult(CVResult):
  72. def _to_str(self, *args, **kwargs):
  73. return super()._to_str(*args, **kwargs).replace("\\\\", "\\")
  74. def _to_img(
  75. self,
  76. ):
  77. """draw formula result"""
  78. try:
  79. env_valid()
  80. except subprocess.CalledProcessError as e:
  81. logging.warning(
  82. "Please refer to 2.3 Formula Recognition Pipeline Visualization in Formula Recognition Pipeline Tutorial to install the LaTeX rendering engine at first."
  83. )
  84. return None
  85. boxes = self["dt_polys"]
  86. formulas = self["rec_formula"]
  87. image = self._img_reader.read(self["input_path"])
  88. h, w = image.height, image.width
  89. img_left = image.copy()
  90. img_right = np.ones((h, w, 3), dtype=np.uint8) * 255
  91. random.seed(0)
  92. draw_left = ImageDraw.Draw(img_left)
  93. if formulas is None or len(formulas) != len(boxes):
  94. formulas = [None] * len(boxes)
  95. for idx, (box, formula) in enumerate(zip(boxes, formulas)):
  96. try:
  97. color = (
  98. random.randint(0, 255),
  99. random.randint(0, 255),
  100. random.randint(0, 255),
  101. )
  102. box = np.array(box)
  103. pts = [(x, y) for x, y in box.tolist()]
  104. draw_left.polygon(pts, outline=color, width=8)
  105. draw_left.polygon(box, fill=color)
  106. img_right_text = draw_box_formula_fine(
  107. (w, h),
  108. box,
  109. formula,
  110. is_debug=False,
  111. )
  112. pts = np.array(box, np.int32).reshape((-1, 1, 2))
  113. cv2.polylines(img_right_text, [pts], True, color, 1)
  114. img_right = cv2.bitwise_and(img_right, img_right_text)
  115. except subprocess.CalledProcessError as e:
  116. logging.warning("Syntax error detected in formula, rendering failed.")
  117. continue
  118. img_left = Image.blend(image, img_left, 0.5)
  119. img_show = Image.new("RGB", (int(w * 2), h), (255, 255, 255))
  120. img_show.paste(img_left, (0, 0, w, h))
  121. img_show.paste(Image.fromarray(img_right), (w, 0, w * 2, h))
  122. return img_show
  123. class FormulaVisualResult(BaseResult):
  124. def __init__(self, data, page_id=None, src_input_name=None):
  125. super().__init__(data)
  126. self.page_id = page_id
  127. self.src_input_name = src_input_name
  128. def _to_str(self, *args, **kwargs):
  129. return super()._to_str(*args, **kwargs).replace("\\\\", "\\")
  130. def get_target_name(self, save_path):
  131. if self.src_input_name.endswith(".pdf"):
  132. save_path = (
  133. Path(save_path)
  134. / f"{Path(self.src_input_name).stem}_pdf"
  135. / Path("page_{:04d}".format(self.page_id + 1))
  136. )
  137. else:
  138. save_path = Path(save_path) / f"{Path(self.src_input_name).stem}"
  139. return save_path
  140. def save_to_json(self, save_path):
  141. if not save_path.lower().endswith(("json")):
  142. save_path = self.get_target_name(save_path)
  143. else:
  144. save_path = Path(save_path).stem
  145. formula_save_path = f"{save_path}_formula.jpg"
  146. self["input_path"] = formula_save_path
  147. self["layout_result"]["input_path"] = formula_save_path
  148. if not str(save_path).endswith(".json"):
  149. save_path = "{}.json".format(save_path)
  150. super().save_to_json(save_path)
  151. def save_to_img(self, save_path):
  152. if not save_path.lower().endswith((".jpg", ".png")):
  153. save_path = self.get_target_name(save_path)
  154. else:
  155. save_path = Path(save_path).stem
  156. formula_save_path = f"{save_path}_formula.jpg"
  157. formula_result = self["formula_result"]
  158. if formula_result:
  159. formula_result.save_to_img(formula_save_path)
  160. def get_align_equation(equation):
  161. is_align = False
  162. equation = str(equation) + "\n"
  163. begin_dict = [
  164. r"begin{align}",
  165. r"begin{align*}",
  166. ]
  167. for begin_sym in begin_dict:
  168. if begin_sym in equation:
  169. is_align = True
  170. break
  171. if not is_align:
  172. equation = (
  173. r"\begin{equation}"
  174. + "\n"
  175. + equation.strip()
  176. + r"\nonumber"
  177. + "\n"
  178. + r"\end{equation}"
  179. + "\n"
  180. )
  181. return equation
  182. def generate_tex_file(tex_file_path, equation):
  183. with open(tex_file_path, "w") as fp:
  184. start_template = (
  185. r"\documentclass{article}" + "\n"
  186. r"\usepackage{cite}" + "\n"
  187. r"\usepackage{amsmath,amssymb,amsfonts}" + "\n"
  188. r"\usepackage{graphicx}" + "\n"
  189. r"\usepackage{textcomp}" + "\n"
  190. r"\DeclareMathSizes{14}{14}{9.8}{7}" + "\n"
  191. r"\pagestyle{empty}" + "\n"
  192. r"\begin{document}" + "\n"
  193. r"\begin{large}" + "\n"
  194. )
  195. fp.write(start_template)
  196. equation = get_align_equation(equation)
  197. fp.write(equation)
  198. end_template = r"\end{large}" + "\n" r"\end{document}" + "\n"
  199. fp.write(end_template)
  200. def generate_pdf_file(tex_path, pdf_dir, is_debug=False):
  201. if os.path.exists(tex_path):
  202. command = "pdflatex -halt-on-error -output-directory={} {}".format(
  203. pdf_dir, tex_path
  204. )
  205. if is_debug:
  206. subprocess.check_call(command, shell=True)
  207. else:
  208. devNull = open(os.devnull, "w")
  209. subprocess.check_call(
  210. command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True
  211. )
  212. def crop_white_area(image):
  213. image = np.array(image).astype("uint8")
  214. gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  215. _, thresh = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)
  216. contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  217. if len(contours) > 0:
  218. x, y, w, h = cv2.boundingRect(np.concatenate(contours))
  219. return [x, y, w, h]
  220. else:
  221. return None
  222. def pdf2img(pdf_path, img_path, is_padding=False):
  223. import fitz
  224. pdfDoc = fitz.open(pdf_path)
  225. if pdfDoc.page_count != 1:
  226. return None
  227. for pg in range(pdfDoc.page_count):
  228. page = pdfDoc[pg]
  229. rotate = int(0)
  230. zoom_x = 2
  231. zoom_y = 2
  232. mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
  233. pix = page.get_pixmap(matrix=mat, alpha=False)
  234. if not os.path.exists(img_path):
  235. os.makedirs(img_path)
  236. pix._writeIMG(img_path, 7, 100)
  237. img = cv2.imread(img_path)
  238. xywh = crop_white_area(img)
  239. if xywh is not None:
  240. x, y, w, h = xywh
  241. img = img[y : y + h, x : x + w]
  242. if is_padding:
  243. img = cv2.copyMakeBorder(
  244. img, 30, 30, 30, 30, cv2.BORDER_CONSTANT, value=(255, 255, 255)
  245. )
  246. return img
  247. return None
  248. def draw_formula_module(img_size, box, formula, is_debug=False):
  249. """draw box formula for module"""
  250. box_width, box_height = img_size
  251. with tempfile.TemporaryDirectory() as td:
  252. tex_file_path = os.path.join(td, "temp.tex")
  253. pdf_file_path = os.path.join(td, "temp.pdf")
  254. img_file_path = os.path.join(td, "temp.jpg")
  255. generate_tex_file(tex_file_path, formula)
  256. if os.path.exists(tex_file_path):
  257. generate_pdf_file(tex_file_path, td, is_debug)
  258. formula_img = None
  259. if os.path.exists(pdf_file_path):
  260. formula_img = pdf2img(pdf_file_path, img_file_path, is_padding=False)
  261. if formula_img is not None:
  262. return formula_img
  263. else:
  264. img_right_text = draw_box_txt_fine(
  265. img_size, box, "Rendering Failed", PINGFANG_FONT_FILE_PATH
  266. )
  267. return img_right_text
  268. def env_valid():
  269. with tempfile.TemporaryDirectory() as td:
  270. tex_file_path = os.path.join(td, "temp.tex")
  271. pdf_file_path = os.path.join(td, "temp.pdf")
  272. img_file_path = os.path.join(td, "temp.jpg")
  273. formula = "a+b=c"
  274. is_debug = False
  275. generate_tex_file(tex_file_path, formula)
  276. if os.path.exists(tex_file_path):
  277. generate_pdf_file(tex_file_path, td, is_debug)
  278. if os.path.exists(pdf_file_path):
  279. formula_img = pdf2img(pdf_file_path, img_file_path, is_padding=False)
  280. def draw_box_formula_fine(img_size, box, formula, is_debug=False):
  281. """draw box formula for pipeline"""
  282. box_height = int(
  283. math.sqrt((box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2)
  284. )
  285. box_width = int(
  286. math.sqrt((box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2)
  287. )
  288. with tempfile.TemporaryDirectory() as td:
  289. tex_file_path = os.path.join(td, "temp.tex")
  290. pdf_file_path = os.path.join(td, "temp.pdf")
  291. img_file_path = os.path.join(td, "temp.jpg")
  292. generate_tex_file(tex_file_path, formula)
  293. if os.path.exists(tex_file_path):
  294. generate_pdf_file(tex_file_path, td, is_debug)
  295. formula_img = None
  296. if os.path.exists(pdf_file_path):
  297. formula_img = pdf2img(pdf_file_path, img_file_path, is_padding=False)
  298. if formula_img is not None:
  299. formula_h, formula_w = formula_img.shape[:-1]
  300. resize_height = box_height
  301. resize_width = formula_w * resize_height / formula_h
  302. formula_img = cv2.resize(
  303. formula_img, (int(resize_width), int(resize_height))
  304. )
  305. formula_h, formula_w = formula_img.shape[:-1]
  306. pts1 = np.float32(
  307. [[0, 0], [box_width, 0], [box_width, box_height], [0, box_height]]
  308. )
  309. pts2 = np.array(box, dtype=np.float32)
  310. M = cv2.getPerspectiveTransform(pts1, pts2)
  311. formula_img = np.array(formula_img, dtype=np.uint8)
  312. img_right_text = cv2.warpPerspective(
  313. formula_img,
  314. M,
  315. img_size,
  316. flags=cv2.INTER_NEAREST,
  317. borderMode=cv2.BORDER_CONSTANT,
  318. borderValue=(255, 255, 255),
  319. )
  320. else:
  321. img_right_text = draw_box_txt_fine(
  322. img_size, box, "Rendering Failed", PINGFANG_FONT_FILE_PATH
  323. )
  324. return img_right_text