Răsfoiți Sursa

add pdf support for formula

liuhongen1234567 1 an în urmă
părinte
comite
40fa4d21d9

+ 10 - 5
paddlex/inference/pipelines/formula_recognition.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 import numpy as np
-from ..components import CropByBoxes
-from ..results import FormulaResult
+from ..components import CropByBoxes, ReadImage
+from ..results import FormulaResult, FormulaVisualResult
 from .base import BasePipeline
 from ...utils import logging
 
@@ -39,6 +39,7 @@ class FormulaRecognitionPipeline(BasePipeline):
             layout_batch_size=layout_batch_size,
             formula_rec_batch_size=formula_rec_batch_size,
         )
+        self.img_reader = ReadImage(format="RGB")
 
     def _build_predictor(self, layout_model, formula_rec_model):
         self.layout_predictor = self._create(model=layout_model)
@@ -56,9 +57,11 @@ class FormulaRecognitionPipeline(BasePipeline):
             self.layout_predictor.set_predictor(device=device)
             self.formula_predictor.set_predictor(device=device)
 
-    def predict(self, x, **kwargs):
+    def predict(self, inputs, **kwargs):
         self.set_predictor(**kwargs)
-        for layout_pred in self.layout_predictor(x):
+        img_info_list = list(self.img_reader(inputs))[0]
+        img_list = [img_info["img"] for img_info in img_info_list]
+        for page_id, layout_pred in enumerate(self.layout_predictor(img_list)):
             single_img_res = {
                 "input_path": "",
                 "layout_result": {},
@@ -89,7 +92,9 @@ class FormulaRecognitionPipeline(BasePipeline):
                         single_img_res["rec_formula"].append(
                             str(formula_res["rec_text"])
                         )
-            yield FormulaResult(single_img_res)
+            single_img_res["formula_result"] = FormulaResult(single_img_res)
+            single_img_res["page_id"] = page_id + 1
+            yield FormulaVisualResult(single_img_res, page_id, inputs)
 
     def sorted_formula_box(self, x):
         coordinate = x["coordinate"]

+ 3 - 2
paddlex/inference/pipelines/layout_parsing/layout_parsing.py

@@ -179,6 +179,7 @@ class LayoutParsingPipeline(_TableRecPipeline):
         for idx, (img_info, layout_pred) in enumerate(
             zip(img_info_list, self.layout_predictor(img_list))
         ):
+            page_id = idx
             single_img_res = {
                 "input_path": "",
                 "layout_result": DetResult({}),
@@ -343,8 +344,8 @@ class LayoutParsingPipeline(_TableRecPipeline):
             single_img_res["ocr_result"] = ocr_res
             single_img_res["table_ocr_result"] = all_table_ocr_res
             single_img_res["layout_parsing_result"] = structure_res
-
-            yield VisualResult(single_img_res)
+            single_img_res["layout_parsing_result"]["page_id"] = page_id + 1
+            yield VisualResult(single_img_res, page_id, inputs)
 
 
 def get_formula_res(predictor, input):

+ 1 - 1
paddlex/inference/results/__init__.py

@@ -21,7 +21,7 @@ from .seal_rec import SealOCRResult
 from .ocr import OCRResult
 from .det import DetResult
 from .seg import SegResult
-from .formula_rec import FormulaRecResult, FormulaResult
+from .formula_rec import FormulaRecResult, FormulaResult, FormulaVisualResult
 from .instance_seg import InstanceSegResult
 from .ts import TSFcResult, TSAdResult, TSClsResult
 from .warp import DocTrResult

+ 57 - 8
paddlex/inference/results/chat_ocr.py

@@ -33,33 +33,82 @@ class VisualInfoResult(BaseResult):
 class VisualResult(BaseResult):
     """VisualInfoResult"""
 
+    def __init__(self, data, page_id=None, src_input_name=None):
+        super().__init__(data)
+        self.page_id = page_id
+        self.src_input_name = src_input_name
+
     def _to_str(self, _, *args, **kwargs):
         return super()._to_str(
             {"layout_parsing_result": self["layout_parsing_result"]}, *args, **kwargs
         )
 
+    def get_target_name(self, save_path):
+        if self.src_input_name.endswith(".pdf"):
+            save_path = (
+                Path(save_path)
+                / f"{Path(self.src_input_name).stem}_pdf"
+                / Path("page_{:04d}".format(self.page_id + 1))
+            )
+        else:
+            save_path = Path(save_path) / f"{Path(self.src_input_name).stem}"
+        return save_path
+
+    def save_to_json(self, save_path):
+        if not save_path.lower().endswith(("json")):
+            save_path = self.get_target_name(save_path)
+        else:
+            save_path = Path(save_path).stem
+        layout_save_path = f"{save_path}_layout.jpg"
+        ocr_save_path = f"{save_path}_ocr.jpg"
+        table_save_path = f"{save_path}_table"
+        self["input_path"] = layout_save_path
+        self["layout_result"]["input_path"] = layout_save_path
+        self["layout_parsing_result"]["input_path"] = layout_save_path
+        self["ocr_result"]["input_path"] = ocr_save_path
+        table_result_num = len(self["table_result"])
+        for idx in range(table_result_num):
+            self["table_result"][idx]["input_path"] = "{}_{:04d}".format(
+                table_save_path, idx + 1
+            )
+
+        if not str(save_path).endswith(".json"):
+            save_path = "{}.json".format(save_path)
+        super().save_to_json(save_path)
+
     def save_to_html(self, save_path):
         if not save_path.lower().endswith(("html")):
-            input_path = self["input_path"]
-            save_path = Path(save_path) / f"{Path(input_path).stem}"
+            save_path = self.get_target_name(save_path)
         else:
             save_path = Path(save_path).stem
-        for table_result in self["table_result"]:
+        table_save_path = f"{save_path}_table"
+        for idx, table_result in enumerate(self["table_result"]):
+            basename = (Path(table_result["input_path"]).name).split(".")[0]
+            table_result["input_path"] = Path(
+                str(table_result["input_path"]).replace(
+                    basename, "{}_{:04d}".format(table_save_path, idx + 1)
+                )
+            )
             table_result.save_to_html(save_path)
 
     def save_to_xlsx(self, save_path):
         if not save_path.lower().endswith(("xlsx")):
-            input_path = self["input_path"]
-            save_path = Path(save_path) / f"{Path(input_path).stem}"
+            save_path = self.get_target_name(save_path)
         else:
             save_path = Path(save_path).stem
-        for table_result in self["table_result"]:
+        table_save_path = f"{save_path}_table"
+        for idx, table_result in enumerate(self["table_result"]):
+            basename = (Path(table_result["input_path"]).name).split(".")[0]
+            table_result["input_path"] = Path(
+                str(table_result["input_path"]).replace(
+                    basename, "{}_{:04d}".format(table_save_path, idx + 1)
+                )
+            )
             table_result.save_to_xlsx(save_path)
 
     def save_to_img(self, save_path):
         if not save_path.lower().endswith((".jpg", ".png")):
-            input_path = self["input_path"]
-            save_path = Path(save_path) / f"{Path(input_path).stem}"
+            save_path = self.get_target_name(save_path)
         else:
             save_path = Path(save_path).stem
 

+ 79 - 8
paddlex/inference/results/formula_rec.py

@@ -20,9 +20,10 @@ import random
 import tempfile
 import subprocess
 import numpy as np
+from pathlib import Path
 from PIL import Image, ImageDraw
 
-from .base import CVResult
+from .base import BaseResult, CVResult
 from ...utils import logging
 from .ocr import draw_box_txt_fine
 from ...utils.fonts import PINGFANG_FONT_FILE_PATH
@@ -36,6 +37,14 @@ class FormulaRecResult(CVResult):
         self,
     ):
         """Draw formula on image"""
+        try:
+            env_valid()
+        except subprocess.CalledProcessError as e:
+            logging.warning(
+                "Please refer to 2.3 Formula Recognition Pipeline Visualization in Formula Recognition Pipeline Tutorial to install the LaTeX rendering engine at first."
+            )
+            return None
+
         image = self._img_reader.read(self["input_path"])
         rec_formula = str(self["rec_text"])
         image = np.array(image.convert("RGB"))
@@ -64,9 +73,7 @@ class FormulaRecResult(CVResult):
             new_image.paste(img_formula, (image.width + 10, 0))
             return new_image
         except subprocess.CalledProcessError as e:
-            logging.warning(
-                "Please refer to 2.3 Formula Recognition Pipeline Visualization in Formula Recognition Pipeline Tutorial to install the LaTeX rendering engine at first."
-            )
+            logging.warning("Syntax error detected in formula, rendering failed.")
             return None
 
 
@@ -79,6 +86,13 @@ class FormulaResult(CVResult):
         self,
     ):
         """draw formula result"""
+        try:
+            env_valid()
+        except subprocess.CalledProcessError as e:
+            logging.warning(
+                "Please refer to 2.3 Formula Recognition Pipeline Visualization in Formula Recognition Pipeline Tutorial to install the LaTeX rendering engine at first."
+            )
+            return None
 
         boxes = self["dt_polys"]
         formulas = self["rec_formula"]
@@ -112,10 +126,8 @@ class FormulaResult(CVResult):
                 cv2.polylines(img_right_text, [pts], True, color, 1)
                 img_right = cv2.bitwise_and(img_right, img_right_text)
             except subprocess.CalledProcessError as e:
-                logging.warning(
-                    "Please refer to 2.3 Formula Recognition Pipeline Visualization in Formula Recognition Pipeline Tutorial to install the LaTeX rendering engine at first."
-                )
-                return None
+                logging.warning("Syntax error detected in formula, rendering failed.")
+                continue
 
         img_left = Image.blend(image, img_left, 0.5)
         img_show = Image.new("RGB", (int(w * 2), h), (255, 255, 255))
@@ -124,6 +136,51 @@ class FormulaResult(CVResult):
         return img_show
 
 
+class FormulaVisualResult(BaseResult):
+
+    def __init__(self, data, page_id=None, src_input_name=None):
+        super().__init__(data)
+        self.page_id = page_id
+        self.src_input_name = src_input_name
+
+    def _to_str(self, *args, **kwargs):
+        return super()._to_str(*args, **kwargs).replace("\\\\", "\\")
+
+    def get_target_name(self, save_path):
+        if self.src_input_name.endswith(".pdf"):
+            save_path = (
+                Path(save_path)
+                / f"{Path(self.src_input_name).stem}_pdf"
+                / Path("page_{:04d}".format(self.page_id + 1))
+            )
+        else:
+            save_path = Path(save_path) / f"{Path(self.src_input_name).stem}"
+        return save_path
+
+    def save_to_json(self, save_path):
+        if not save_path.lower().endswith(("json")):
+            save_path = self.get_target_name(save_path)
+        else:
+            save_path = Path(save_path).stem
+
+        formula_save_path = f"{save_path}_formula.jpg"
+        self["input_path"] = formula_save_path
+        self["layout_result"]["input_path"] = formula_save_path
+        if not str(save_path).endswith(".json"):
+            save_path = "{}.json".format(save_path)
+        super().save_to_json(save_path)
+
+    def save_to_img(self, save_path):
+        if not save_path.lower().endswith((".jpg", ".png")):
+            save_path = self.get_target_name(save_path)
+        else:
+            save_path = Path(save_path).stem
+        formula_save_path = f"{save_path}_formula.jpg"
+        formula_result = self["formula_result"]
+        if formula_result:
+            formula_result.save_to_img(formula_save_path)
+
+
 def get_align_equation(equation):
     is_align = False
     equation = str(equation) + "\n"
@@ -247,6 +304,20 @@ def draw_formula_module(img_size, box, formula, is_debug=False):
         return img_right_text
 
 
+def env_valid():
+    with tempfile.TemporaryDirectory() as td:
+        tex_file_path = os.path.join(td, "temp.tex")
+        pdf_file_path = os.path.join(td, "temp.pdf")
+        img_file_path = os.path.join(td, "temp.jpg")
+        formula = "a+b=c"
+        is_debug = False
+        generate_tex_file(tex_file_path, formula)
+        if os.path.exists(tex_file_path):
+            generate_pdf_file(tex_file_path, td, is_debug)
+        if os.path.exists(pdf_file_path):
+            formula_img = pdf2img(pdf_file_path, img_file_path, is_padding=False)
+
+
 def draw_box_formula_fine(img_size, box, formula, is_debug=False):
     """draw box formula for pipeline"""
     box_height = int(