rewrite with pypdf and ocrmypdf

app.py

@@ -1,88 +1,39 @@
 import spaces
 import gradio as gr
-import surya.detection as detection
-import surya.layout as layout
-import os
-import base64
-
-
-# Monkey patch to prevent spawning processes
-def batch_text_detection(images, model, processor, batch_size=None):
-    preds, orig_sizes = detection.batch_detection(
-        images, model, processor, batch_size=batch_size
-    )
-    results = []
-    for i in range(len(images)):
-        result = detection.parallel_get_lines(preds[i], orig_sizes[i])
-        results.append(result)
-
-    return results
-
-
-detection.batch_text_detection = batch_text_detection
-
-
-def batch_layout_detection(
-    images, model, processor, detection_results=None, batch_size=None
-):
-    preds, orig_sizes = layout.batch_detection(
-        images, model, processor, batch_size=batch_size
-    )
-    id2label = model.config.id2label
-
-    results = []
-    for i in range(len(images)):
-        result = layout.parallel_get_regions(
-            preds[i],
-            orig_sizes[i],
-            id2label,
-            detection_results[i] if detection_results else None,
-        )
-        results.append(result)
-
-    return results
-
-
-layout.batch_layout_detection = batch_layout_detection
-
-from marker.convert import convert_single_pdf
-from marker.models import load_all_models
-
-model_list = load_all_models()
+from pypdf import PdfReader
+import ocrmypdf
 
 
 @spaces.GPU
-def convert(pdf_file, extract_images):
-    global model_list
-
-    full_text, images, out_meta = convert_single_pdf(
-        pdf_file, model_list, batch_multiplier=16
-    )
-    image_data = {}
-    if extract_images:
-        for filename, image in images.items():
-            image.save(filename, "PNG")
+def convert(pdf_file):
+    reader = PdfReader(pdf_file)
 
-            with open(filename, "rb") as f:
-                image_bytes = f.read()
+    # Check if there are any images
+    image_count = 0
+    for page in reader.pages:
+        image_count += len(page.images)
 
-            image_base64 = base64.b64encode(image_bytes).decode("utf-8")
-            image_data[filename] = image_base64
+    # If there are images, perform OCR on the document
+    if image_count > 0:
+        out_pdf_file = pdf_file.replace(".pdf", "_ocr.pdf")
+        ocrmypdf.ocr(pdf_file, out_pdf_file, force_ocr=True)
+        pdf_file = out_pdf_file
 
-            os.remove(filename)
+    # Extract text
+    full_text = ""
+    for idx, page in enumerate(reader.pages):
+        full_text += f"\n\n---- Page {idx} ----\n\n" + page.extract_text()
 
-    return full_text, out_meta, image_data
+    return full_text, reader.metadata
 
 
 gr.Interface(
     convert,
     inputs=[
         gr.File(label="Upload PDF", type="filepath"),
-        gr.Checkbox(label="Extract Images"),
     ],
     outputs=[
         gr.Text(label="Markdown"),
         gr.JSON(label="Metadata"),
-        gr.JSON(label="Images"),
     ],
 ).launch()

packages.txt

@@ -0,0 +1,2 @@
+ocrmypdf
+tesseract-ocr-eng

requirements.txt

@@ -1 +1,2 @@
-marker-pdf==0.2.5
+ocrmypdf==16.3.1
+pypdf==4.2.0