extract text and tables

app.py

@@ -1,34 +1,45 @@
 from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import JSONResponse
-import pypdfium2 as pdfium
 import base64
 import re
 
 app = FastAPI()
 
-def extract_text_from_pdf(pdf_bytes: bytes) -> str:
+from io import BytesIO
+from gmft.pdf_bindings import PyPDFium2Document
+from utils import get_page_text_with_tables, detector, formatter
+
+def extract_text_from_pdf(pdf_bytes: bytes, page_numbers=None) -> str:
     """
-    Extract text from PDF bytes using pypdfium2
+    Extract text from PDF bytes using gmft without temporary files
     """
-    pdf_file = pdfium.PdfDocument(pdf_bytes)
-    text_parts = []
+    # Create a PyPDFium2Document directly from bytes
+    doc = PyPDFium2Document(pdf_bytes)
+    page_set = set(page_numbers if page_numbers else list(range(len(doc))))
     
     try:
-        for page in pdf_file:
-            textpage = page.get_textpage()
-            text = textpage.get_text_range()
-            text_parts.append(text)
+        pages = []
+        for page_num, page in enumerate(doc):
+            if not page_num in page_set:
+                continue
+            try:
+                tables = detector.extract(page)
+                fmt_tables = [formatter.extract(table, margin=(0, 0, 0, 0)) for table in tables]
+                page_text = get_page_text_with_tables(page, fmt_tables)
+                pages.append(page_text)
+            finally:
+                page.close()
     finally:
-        pdf_file.close()
+        doc.close()
     
-    return "\n".join(text_parts)
+    return pages
 
 @app.get("/")
 def greet_json():
     return {"Hello": "World!"}
 
 @app.post("/extract-text")
-async def extract_pdf_text(file: UploadFile = File(...)):
+async def extract_pdf_text(file: UploadFile = File(...), page_numbers: str = None):
     """
     Endpoint to extract text from uploaded PDF file
     """
@@ -42,9 +53,21 @@ async def extract_pdf_text(file: UploadFile = File(...)):
     # Read the file content
     content = await file.read()
     
+    # Parse page_numbers if provided
+    parsed_page_numbers = None
+    if page_numbers:
+        try:
+            # Convert comma-separated string to list of integers
+            parsed_page_numbers = [int(p.strip()) for p in page_numbers.split(',') if p.strip()]
+        except ValueError:
+            return JSONResponse(
+                status_code=400,
+                content={"error": "Invalid page_numbers format. Use comma-separated integers."}
+            )
+    
     try:
         # Extract text from PDF
-        extracted_text = extract_text_from_pdf(content)
+        extracted_text = extract_text_from_pdf(content, parsed_page_numbers)
         
         return {
             "filename": file.filename,
@@ -74,6 +97,27 @@ async def extract_pdf_text_base64(data: dict):
     # Extract filename if provided
     filename = data.get('filename', 'unknown.pdf')
     
+    # Extract page_numbers if provided
+    page_numbers = data.get('page_numbers')
+    parsed_page_numbers = None
+    if page_numbers:
+        try:
+            # Handle both string and list formats
+            if isinstance(page_numbers, str):
+                parsed_page_numbers = [int(p.strip()) for p in page_numbers.split(',') if p.strip()]
+            elif isinstance(page_numbers, list):
+                parsed_page_numbers = [int(p) for p in page_numbers if isinstance(p, (int, str))]
+            else:
+                return JSONResponse(
+                    status_code=400,
+                    content={"error": "Invalid page_numbers format. Use comma-separated integers or array."}
+                )
+        except (ValueError, TypeError):
+            return JSONResponse(
+                status_code=400,
+                content={"error": "Invalid page_numbers format. Use comma-separated integers or array."}
+            )
+    
     try:
         # Handle data URL format (e.g., "data:application/pdf;base64,...")
         if base64_string.startswith('data:'):
@@ -90,7 +134,7 @@ async def extract_pdf_text_base64(data: dict):
         pdf_bytes = base64.b64decode(base64_string)
         
         # Extract text from PDF
-        extracted_text = extract_text_from_pdf(pdf_bytes)
+        extracted_text = extract_text_from_pdf(pdf_bytes, parsed_page_numbers)
         
         return {
             "filename": filename,

requirements.txt

@@ -1,4 +1,5 @@
 fastapi[all]
 aiohttp
 uvicorn[standard]
-pypdfium2
+pypdfium2
+gmft

test_pdf_endpoint.py

@@ -6,10 +6,11 @@ import base64
 async def test_pdf_extraction():
     # Check if filename is provided as argument
     if len(sys.argv) < 2:
-        print("Usage: python test_pdf_endpoint.py <pdf_filename>")
+        print("Usage: python test_pdf_endpoint.py <pdf_filename> [page_numbers]")
         return
     
     pdf_filename = sys.argv[1]
+    page_numbers = sys.argv[2] if len(sys.argv) > 2 else None
     
     # Read the PDF file
     try:
@@ -25,6 +26,8 @@ async def test_pdf_extraction():
     # Test regular file upload endpoint
     print("\n--- Testing file upload endpoint ---")
     url = "http://localhost:8000/extract-text"
+    if page_numbers:
+        url += f"?page_numbers={page_numbers}"
     
     try:
         async with aiohttp.ClientSession() as session:
@@ -61,6 +64,8 @@ async def test_pdf_extraction():
             "file": base64_string,
             "filename": pdf_filename
         }
+        if page_numbers:
+            payload["page_numbers"] = page_numbers
         
         async with aiohttp.ClientSession() as session:
             async with session.post(

utils.py

@@ -0,0 +1,52 @@
+from gmft.formatters.base import FormattedTable
+from gmft.formatters.page.base import FormattedPage
+from gmft.pdf_bindings.base import BasePage
+
+from gmft.auto import AutoTableFormatter, AutoTableDetector, TATRFormatConfig
+from gmft.pdf_bindings import PyPDFium2Document
+
+detector = AutoTableDetector()
+formatter = AutoTableFormatter(TATRFormatConfig(semantic_spanning_cells = False, enable_multi_header=False))
+
+
+def get_page_text_with_tables(
+    page: BasePage, tables: list[FormattedTable]
+) -> FormattedPage:
+    if not tables:
+        return page._get_text_with_breaks()
+
+    text_builder = []
+    done = [False for _ in tables]
+    for (
+        x0,
+        y0,
+        x1,
+        y1,
+        word,
+        blockno,
+        lineno,
+        wordno,
+    ) in page._get_positions_and_text_and_breaks():
+        for j, table in enumerate(tables):
+            if table.rect.is_intersecting((x0, y0, x1, y1)):
+                if not done[j]:
+                    try:
+                      table_content = table.df().fillna("").to_latex(index=False)
+                      text_builder.append(f"\n{table_content}\n")
+                    except:
+                      # it throws errors when tables have no text
+                      pass
+                    done[j] = True
+                break
+        else:
+            # no table found
+            if wordno == 0:
+                text_builder.append("\n")
+                if lineno == 0:
+                  text_builder.append("\n")
+            else:
+                text_builder.append(" ")
+            text_builder.append(word)
+
+    page_content = "".join(text_builder).lstrip()
+    return page_content