formula_rec.py 11 KB

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