feat: add img2table gmft

app.py

@@ -1,3 +1,7 @@
+import os
+import zipfile
+from collections import defaultdict
+
 from utils import fix_problematic_imports, prepare_env_mineru
 
 fix_problematic_imports()  # noqa
@@ -13,18 +17,23 @@ from gradio_pdf import PDF
 
 from backends import (
     convert_docling,
+    convert_gemini,
+    convert_gmft,
+    convert_img2table,
     convert_marker,
     convert_mineru,
     convert_unstructured,
+    convert_zerox,
 )
 from backends.settings import ENABLE_DEBUG_MODE
 from utils import remove_images_from_markdown, trim_pages
 
 TRIMMED_PDF_PATH = Path("/tmp/trimmed_input")
 TRIMMED_PDF_PATH.mkdir(exist_ok=True)
+DO_WARMUP = os.getenv("DO_WARMUP", "True").lower() == "true"
 
 
-def convert_document(path, method, enabled=True):
+def convert_document(path, method, start_page=0, enabled=True):
     if enabled:
         print("Processing file", path, "with method", method)
     else:
@@ -33,7 +42,11 @@ def convert_document(path, method, enabled=True):
     # benchmarking
     start = time.time()
 
-    path = trim_pages(path, output_path=TRIMMED_PDF_PATH)
+    path = trim_pages(
+        path,
+        output_path=TRIMMED_PDF_PATH,
+        start_page=start_page,
+    )
     file_name = Path(path).stem
     debug_image_paths = []
     text = "unknown method"
@@ -51,6 +64,16 @@ def convert_document(path, method, enabled=True):
         )
     elif method == "MinerU":
         text, debug_image_paths = convert_mineru(path, file_name)
+    elif method == "Gemini (API)":
+        text, debug_image_paths = convert_gemini(path, file_name)
+    elif method == "Zerox":
+        text, debug_image_paths = convert_zerox(path, file_name)
+    elif method == "Img2Table":
+        text, debug_image_paths = convert_img2table(path, file_name)
+    elif method == "GMFT":
+        text, debug_image_paths = convert_gmft(path, file_name)
+    else:
+        raise ValueError(f"Unsupported method: {method}")
 
     duration = time.time() - start
     duration_message = f"Conversion with {method} took *{duration:.2f} seconds*"
@@ -63,6 +86,51 @@ def convert_document(path, method, enabled=True):
     )
 
 
+def to_zip_file(file_path, methods, *output_components):
+    markdown_text_dict = dict()
+    debug_images_dict = defaultdict(list)
+    for idx, method_name in enumerate(SUPPORTED_METHODS):
+        if method_name not in methods:
+            continue
+
+        markdown_text = output_components[idx * 4 + 2]
+        debug_images = output_components[idx * 4 + 3]
+
+        markdown_text_dict[method_name] = markdown_text
+        debug_images_dict[method_name] = debug_images
+
+    # create new temp directory using Python's tempfile module
+    temp_dir = Path(file_path).parent
+    zip_file_path = temp_dir / "output.zip"
+
+    markdown_path = temp_dir / f"{method_name}.md"
+    with open(markdown_path, "w") as f:
+        f.write(markdown_text)
+
+    # create a zip file in write mode
+    with zipfile.ZipFile(zip_file_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for method_name, markdown_text in markdown_text_dict.items():
+            debug_image_paths = debug_images_dict[method_name]
+
+            # write the markdown text to the zip file
+            zipf.write(
+                markdown_path,
+                f"{method_name}/{method_name}.md",
+            )
+            if debug_image_paths:
+                for idx, (debug_image_path, _) in enumerate(debug_image_paths):
+                    debug_image_name = Path(debug_image_path).name
+                    zipf.write(
+                        debug_image_path,
+                        f"{method_name}/{debug_image_name}",
+                    )
+
+    return gr.update(
+        value=str(zip_file_path),
+        visible=True,
+    )
+
+
 def show_tabs(selected_methods):
     visible_tabs = []
     for method in SUPPORTED_METHODS:
@@ -79,14 +147,25 @@ latex_delimiters = [
 # startup test (also for loading models the first time)
 start_startup = time.time()
 WARMUP_PDF_PATH = "table.pdf"
-SUPPORTED_METHODS = ["PyMuPDF", "Docling", "Marker", "MinerU", "Unstructured"]
+SUPPORTED_METHODS = [
+    "PyMuPDF",
+    "Docling",
+    "Marker",
+    "MinerU",
+    "Unstructured",
+    "Gemini (API)",
+    "Img2Table",
+    "GMFT",
+    # "Zerox"
+]
 
-print("Warm-up sequence")
-for method in SUPPORTED_METHODS:
-    for _ in range(1):
-        convert_document(WARMUP_PDF_PATH, method)
-startup_duration = time.time() - start_startup
-print(f"Total start-up time: {startup_duration:.2f} seconds")
+if DO_WARMUP:
+    print("Warm-up sequence")
+    for method in SUPPORTED_METHODS:
+        for _ in range(1):
+            convert_document(WARMUP_PDF_PATH, method)
+    startup_duration = time.time() - start_startup
+    print(f"Total start-up time: {startup_duration:.2f} seconds")
 
 with gr.Blocks(
     theme=gr.themes.Ocean(),
@@ -106,7 +185,28 @@ with gr.Blocks(
                     ".pdf",
                 ],
             )
+            with gr.Accordion(
+                "Advanced settings",
+                open=False,
+            ):
+                start_page = gr.Number(
+                    label="Starting page (only max 5 consecutive pages are processed)",
+                    minimum=1,
+                    maximum=100,
+                    step=1,
+                    value=1,
+                )
+                visual_checkbox = gr.Checkbox(
+                    label="Enable debug visualization",
+                    visible=ENABLE_DEBUG_MODE,
+                    value=True,
+                )
             progress_status = gr.Markdown("", show_label=False, container=False)
+            output_file = gr.File(
+                label="Download output",
+                interactive=False,
+                visible=False,
+            )
 
         with gr.Column(variant="panel", scale=5):
             with gr.Row():
@@ -116,12 +216,6 @@ with gr.Blocks(
                     value=SUPPORTED_METHODS[:2],
                     multiselect=True,
                 )
-            with gr.Row():
-                visual_checkbox = gr.Checkbox(
-                    label="Enable debug visualization",
-                    visible=ENABLE_DEBUG_MODE,
-                    value=True,
-                )
             with gr.Row():
                 convert_btn = gr.Button("Convert", variant="primary", scale=2)
                 clear_btn = gr.ClearButton(value="Clear", scale=1)
@@ -210,11 +304,14 @@ with gr.Blocks(
 
             return msg
 
-        def process_method(input_file, selected_methods, method=method):
+        def process_method(input_file, start_page, selected_methods, method=method):
             if input_file is None:
                 raise ValueError("Please upload a PDF file first!")
             return convert_document(
-                input_file, method=method, enabled=method in selected_methods
+                input_file,
+                method=method,
+                start_page=start_page - 1,
+                enabled=method in selected_methods,
             )
 
         click_event = click_event.then(
@@ -222,25 +319,35 @@ with gr.Blocks(
             inputs=[methods],
             outputs=[progress_status],
         ).then(
-            fn=lambda input_file, methods, method=method: process_method(
-                input_file, methods, method
+            fn=lambda input_file, start_page, methods, method=method: process_method(
+                input_file, start_page, methods, method
             ),
-            inputs=[input_file, methods],
+            inputs=[input_file, start_page, methods],
             outputs=output_components[idx * 4 : (idx + 1) * 4],
         )
 
-    click_event.then(
-        lambda: "All tasks completed.",
-        outputs=[progress_status],
+    click_event.then(lambda: "All tasks completed.", outputs=[progress_status],).then(
+        fn=to_zip_file,
+        inputs=[
+            input_file,
+            methods,
+        ]
+        + output_components,
+        outputs=[output_file],
     )
 
     clear_btn.add(
         [
             input_file,
             pdf_preview,
+            output_file,
         ]
         + output_components
     )
+    clear_btn.click(
+        fn=lambda: gr.update(visible=False),
+        outputs=[output_file],
+    )
 
     visual_checkbox.change(
         fn=lambda state: [gr.update(visible=state)] * len(visualization_sub_tabs),
@@ -248,4 +355,7 @@ with gr.Blocks(
         outputs=visualization_sub_tabs,
     )
 
-    demo.launch(show_error=True)
+    demo.launch(
+        show_error=True,
+        max_file_size="50mb",
+    )

backends/__init__.py

@@ -1,11 +1,19 @@
 from .docling import convert_docling
+from .gemini import convert_gemini
+from .gmft import convert_gmft
+from .img2table import convert_img2table
 from .marker import convert_marker
 from .mineru import convert_mineru
 from .unstructured import convert_unstructured
+from .zerox import convert_zerox
 
 __all__ = [
     "convert_docling",
     "convert_marker",
     "convert_mineru",
     "convert_unstructured",
+    "convert_gemini",
+    "convert_zerox",
+    "convert_img2table",
+    "convert_gmft",
 ]

backends/gemini.py

@@ -0,0 +1,43 @@
+import os
+from pathlib import Path
+
+from google import genai
+from google.genai import types
+
+# Create a client
+client = genai.Client(api_key=os.getenv("GEMINI_API_KEY", ""))
+MODEL_NAME = "gemini-2.0-flash"
+PROMPT = """
+Convert the following document to markdown, preserving header, table and figure structure as much as possible.
+Return only the markdown with no explanation text. Do not include delimiters like ```markdown or ```html.
+
+RULES:
+    - You must include all information on the page. Do not exclude headers, footers, or subtext.
+    - Return tables in Markdown format.
+    - Must format headers / sub-headers in Markdown format (#, ##, etc).
+    - Attempt to merge line-breaks in to coherent paragraphs.
+    - Charts & infographics must be interpreted to a text-based markdown format. Prefer table format when applicable.
+    - Do not include any images URL / tag in the markdown.
+    - Page numbers should be wrapped in brackets. Ex: <page_number>14<page_number> or <page_number>9/22<page_number>
+    - Prefer using ☐ and ☑ for check boxes.
+"""  # noqa: E501
+
+
+def convert_gemini(path: str, file_name: str):
+    # Generate a structured response using the Gemini API
+    generation_config = types.GenerationConfig(
+        max_output_tokens=8192,
+    ).to_json_dict()
+    response = client.models.generate_content(
+        model=MODEL_NAME,
+        contents=[
+            PROMPT,
+            types.Part.from_bytes(
+                data=Path(path).read_bytes(),
+                mime_type="application/pdf",
+            ),
+        ],
+        config=generation_config,
+    )
+    # Convert the response to the pydantic model and return it
+    return response.text, []

backends/gmft.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+
+from gmft.auto import AutoFormatConfig, AutoTableFormatter, CroppedTable, TableDetector
+from gmft.pdf_bindings import PyPDFium2Document
+
+from .settings import ENABLE_DEBUG_MODE
+
+detector = TableDetector()
+config = AutoFormatConfig()
+config.semantic_spanning_cells = True  # [Experimental] better spanning cells
+config.enable_multi_header = True  # multi-headers
+formatter = AutoTableFormatter(config)
+
+
+GMFT_DEBUG_PATH = Path("/tmp/gmft")
+GMFT_DEBUG_PATH.mkdir(exist_ok=True)
+
+
+def ingest_pdf(pdf_path) -> list[CroppedTable]:
+    doc = PyPDFium2Document(pdf_path)
+
+    tables = []
+    for page in doc:
+        tables += detector.extract(page)
+    return tables
+
+
+def convert_gmft(path: str, file_name: str):
+    tables = ingest_pdf(path)
+    formatted_tables = []
+    debug_image_paths = []
+
+    debug_path = GMFT_DEBUG_PATH / file_name
+    debug_path.mkdir(exist_ok=True)
+
+    for idx, table in enumerate(tables):
+        ft = formatter.extract(
+            table,
+            dpi=72 * 2,
+        )
+        df = ft.df()
+        if df is not None:
+            html = df.fillna("").to_html(
+                index=False,
+            )
+            formatted_tables.append(html)
+
+        if ENABLE_DEBUG_MODE:
+            image_path = debug_path / f"table_{idx}.png"
+            ft.image().save(image_path)
+            debug_image_paths.append(image_path)
+
+    content = "\n\n".join(formatted_tables)
+    return content, debug_image_paths

backends/img2table.py

@@ -0,0 +1,55 @@
+from pathlib import Path
+
+import cv2
+from img2table.document import PDF
+from img2table.ocr import SuryaOCR
+
+from .settings import ENABLE_DEBUG_MODE
+
+ocr = SuryaOCR(
+    langs=["en"],
+)
+IMG2TABLE_DEBUG_PATH = Path("/tmp/img2table")
+IMG2TABLE_DEBUG_PATH.mkdir(exist_ok=True)
+
+
+def convert_img2table(path: str, file_name: str):
+    doc = PDF(path)
+    pages = doc.extract_tables(
+        ocr=ocr,
+        implicit_rows=False,
+        implicit_columns=False,
+        borderless_tables=True,
+        min_confidence=50,
+    )
+    debug_image_paths = []
+
+    if ENABLE_DEBUG_MODE:
+        debug_path = IMG2TABLE_DEBUG_PATH / file_name
+        debug_path.mkdir(exist_ok=True)
+
+        images = doc.images
+        for idx, page_number in enumerate(doc.pages or range(len(images))):
+            page_image = images[idx]
+            for table in pages[page_number]:
+                for row in table.content.values():
+                    for cell in row:
+                        cv2.rectangle(
+                            page_image,
+                            (cell.bbox.x1, cell.bbox.y1),
+                            (cell.bbox.x2, cell.bbox.y2),
+                            (0, 0, 255),
+                            2,
+                        )
+            image_path = debug_path / f"page_{idx}.png"
+            debug_image_paths.append(image_path)
+            cv2.imwrite(str(image_path), page_image)
+
+    content = "\n\n".join(
+        [
+            (table.title if table.title else "") + "\n\n" + table.html
+            for tables in pages.values()
+            for table in tables
+        ]
+    )
+    return content, debug_image_paths

backends/zerox.py

@@ -0,0 +1,43 @@
+import asyncio
+import re
+from pathlib import Path
+
+from pyzerox import zerox
+
+
+def remove_images_from_markdown(markdown_text):
+    # remove <image> and ![image](path) from markdown
+    markdown_text = re.sub(r"<img[^>]*>", "", markdown_text)
+    markdown_text = re.sub(r"!\[[^\]]*\]\([^)]*\)", "", markdown_text)
+    return markdown_text
+
+
+ZEROX_DEBUG_PATH = Path("/tmp/zerox_debug")
+ZEROX_DEBUG_PATH.mkdir(exist_ok=True)
+MODEL_NAME = "gemini/gemini-2.0-flash"
+
+
+def clean_up_html_code_block(text: str):
+    # remove ```html and ``` from text
+    text = text.replace("```html", "")
+    text = text.replace("```", "")
+    return text
+
+
+def convert_zerox(path: str, file_name: str):
+    output_dir = ZEROX_DEBUG_PATH / file_name
+    output_dir.mkdir(exist_ok=True)
+
+    async def async_convert():
+        return await zerox(
+            concurrency=4,
+            file_path=path,
+            model=MODEL_NAME,
+            output_dir=output_dir,
+        )
+
+    output = asyncio.run(async_convert())
+    output_text = "\n\n".join(page.content for page in output.pages)
+    output_text = clean_up_html_code_block(output_text)
+    output_text = remove_images_from_markdown(output_text)
+    return output_text, []

requirements.txt

@@ -14,5 +14,7 @@ unstructured[pdf]
 ultralytics>=8.3.48
 transformers<5.0.0,>=4.45.2
 struct-eqtable==0.3.2
-openai
 doclayout_yolo==0.0.2b1
+openai
+opencv-contrib-python
+gmft

utils.py

@@ -14,14 +14,14 @@ def remove_images_from_markdown(markdown_text):
 
 
 @functools.lru_cache(maxsize=None)
-def trim_pages(pdf_path, output_path, trim_pages=5):
+def trim_pages(pdf_path, output_path, start_page=0, trim_pages=5):
     doc = pymupdf.open(pdf_path)
     parent_dir_name = Path(pdf_path).parent.name
     output_file_path = Path(output_path) / f"{parent_dir_name}.pdf"
 
     num_pages = len(doc)
     if num_pages > trim_pages:
-        to_select = list(range(trim_pages))
+        to_select = list(range(start_page, min(start_page + trim_pages, num_pages)))
         doc.select(to_select)
         doc.ez_save(output_file_path)
         print("Trimmed pdf to with pages", to_select, "path", output_file_path)