From 6c4aa1316a6c059b298065ab3eb5646ca32d27b7 Mon Sep 17 00:00:00 2001
From: tiankuolu <tiankuolu154@gmail.com>
Date: Wed, 29 Jul 2026 21:58:53 +0800
Subject: [PATCH] feat: add PDF batch processing and merge results to prevent
 memory overflow

---
 .gitignore       |  1 +
 README.md        | 16 +++++++++++
 merge_results.py | 41 ++++++++++++++++++++++++++++
 run_ocr.py       | 69 ++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 127 insertions(+)
 create mode 100644 merge_results.py
 create mode 100644 run_ocr.py

diff --git a/.gitignore b/.gitignore
index d3d3fcc..85e0481 100644
--- a/.gitignore
+++ b/.gitignore
@@ -27,6 +27,7 @@ htmlcov/
 
 # Local runtime outputs
 outputs/
+ocr_output/
 log/
 *.log
 
diff --git a/README.md b/README.md
index afd12ad..61a274c 100644
--- a/README.md
+++ b/README.md
@@ -130,6 +130,22 @@ model.infer_multi(
 )
 ```
 
+#### Run and merge a PDF page by page
+
+Run OCR:
+
+```shell
+python run_ocr.py your_doc.pdf
+```
+
+Merge the page results:
+
+```shell
+python merge_results.py ./ocr_output
+```
+
+The merged Markdown file is saved to `./ocr_output/merged.md`.
+
 ### vLLM
 
 Please refer to the official vLLM recipe for deployment details:
diff --git a/merge_results.py b/merge_results.py
new file mode 100644
index 0000000..79e4628
--- /dev/null
+++ b/merge_results.py
@@ -0,0 +1,41 @@
+import os
+import shutil
+import sys
+
+output_dir = sys.argv[1] if len(sys.argv) > 1 else './ocr_output'
+images_dir = os.path.join(output_dir, 'images')
+os.makedirs(images_dir, exist_ok=True)
+
+merged_path = os.path.join(output_dir, 'merged.md')
+page_count = 0
+with open(merged_path, 'w', encoding='utf-8') as merged:
+    for name in sorted(os.listdir(output_dir)):
+        page_dir = os.path.join(output_dir, name)
+        if not os.path.isdir(page_dir) or not name.startswith('page_'):
+            continue
+
+        md_path = os.path.join(page_dir, 'result.md')
+        if not os.path.exists(md_path):
+            continue
+
+        with open(md_path, 'r', encoding='utf-8') as page_file:
+            text = page_file.read()
+
+        src_images = os.path.join(page_dir, 'images')
+        if os.path.isdir(src_images):
+            for img_name in sorted(os.listdir(src_images)):
+                src = os.path.join(src_images, img_name)
+                dst_name = f'{name}_{img_name}'
+                shutil.copy2(src, os.path.join(images_dir, dst_name))
+                text = text.replace(
+                    f'images/{img_name}',
+                    f'images/{dst_name}',
+                )
+
+        if page_count:
+            merged.write('\n\n---\n\n')
+        merged.write(text.strip())
+        page_count += 1
+
+print(f"Merged {page_count} pages → {merged_path}")
+print(f"Images → {images_dir}/")
diff --git a/run_ocr.py b/run_ocr.py
new file mode 100644
index 0000000..539a60d
--- /dev/null
+++ b/run_ocr.py
@@ -0,0 +1,69 @@
+import os
+import sys
+
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+model_name = 'baidu/Unlimited-OCR'
+
+print("Loading...")
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    model_name, trust_remote_code=True,
+    use_safetensors=True, torch_dtype=torch.bfloat16,
+)
+model = model.eval().cuda()
+print("Ready.\n")
+
+if len(sys.argv) < 2:
+    print("Usage: python run_ocr.py <image_or_pdf_path>")
+    sys.exit(0)
+
+input_path = sys.argv[1]
+ext = os.path.splitext(input_path)[1].lower()
+output_dir = './ocr_output'
+os.makedirs(output_dir, exist_ok=True)
+
+if ext == '.pdf':
+    import fitz
+
+    doc = fitz.open(input_path)
+    total = len(doc)
+    print(f"PDF: {total} pages")
+
+    completed = 0
+    for i in range(total):
+        page = doc[i]
+        pix = page.get_pixmap(matrix=fitz.Matrix(300 / 72, 300 / 72))
+        tmp_path = os.path.join(output_dir, f'_page_{i+1:04d}.png')
+        pix.save(tmp_path)
+
+        print(f"  Page {i+1}/{total}...", end=" ", flush=True)
+        try:
+            model.infer(tokenizer,
+                prompt='<image>document parsing.',
+                image_file=tmp_path,
+                output_path=os.path.join(output_dir, f'page_{i+1:04d}'),
+                base_size=1024, image_size=1024, crop_mode=False,
+                max_length=32768, no_repeat_ngram_size=35, ngram_window=128,
+                save_results=True,
+            )
+            completed += 1
+            print("OK")
+        except Exception as e:
+            print(f"FAILED: {e}")
+        finally:
+            os.remove(tmp_path)
+
+    doc.close()
+    print(f"\nDone. {completed}/{total} pages")
+else:
+    print(f"Image: {input_path}")
+    model.infer(tokenizer,
+        prompt='<image>document parsing.', image_file=input_path,
+        output_path=output_dir,
+        base_size=1024, image_size=640, crop_mode=True,
+        max_length=32768, no_repeat_ngram_size=35, ngram_window=128,
+        save_results=True,
+    )
+    print(f"Done. Output: {output_dir}")