testing out OCR viewer of OCR'd text
This commit is contained in:
nathan
220
scripts/ocr_wardiaries.py
Normal file
220
scripts/ocr_wardiaries.py
Normal file
@@ -0,0 +1,220 @@
|
||||
"""
|
||||
ocr_wardiaries.py
|
||||
-----------------
|
||||
Converts war diary PDFs to Markdown using the olmOCR model hosted on DeepInfra.
|
||||
Bypasses the olmOCR pipeline GPU check by calling the API directly.
|
||||
|
||||
Requirements:
|
||||
pip install requests pillow
|
||||
|
||||
Poppler must be on PATH (pdftoppm command must work).
|
||||
|
||||
Usage:
|
||||
python ocr_wardiaries.py --input_dir "C:\path\to\pdfs" --output_dir "C:\path\to\output" --api_key "YOUR_KEY"
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
DEEPINFRA_API_URL = "https://api.deepinfra.com/v1/openai/chat/completions"
|
||||
MODEL = "allenai/olmOCR-2-7B-1025"
|
||||
DPI = 150 # Higher = better quality but slower/more expensive. 150 is good for typewritten text.
|
||||
MAX_RETRIES = 3 # Retries per page on API error
|
||||
RETRY_DELAY = 5 # Seconds between retries
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def pdf_to_images(pdf_path: Path, output_dir: Path, dpi: int = DPI) -> list[Path]:
|
||||
"""Render each page of a PDF to a PNG image using pdftoppm (poppler)."""
|
||||
prefix = output_dir / pdf_path.stem
|
||||
cmd = [
|
||||
"pdftoppm",
|
||||
"-r", str(dpi),
|
||||
"-png",
|
||||
str(pdf_path),
|
||||
str(prefix),
|
||||
]
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"pdftoppm failed for {pdf_path.name}:\n{result.stderr}")
|
||||
|
||||
# pdftoppm names files like prefix-1.png, prefix-2.png etc (zero-padded)
|
||||
images = sorted(output_dir.glob(f"{pdf_path.stem}-*.png"))
|
||||
return images
|
||||
|
||||
|
||||
def image_to_base64(image_path: Path) -> str:
|
||||
"""Convert an image file to a base64 string."""
|
||||
with open(image_path, "rb") as f:
|
||||
return base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
|
||||
def ocr_page(image_path: Path, api_key: str, page_num: int) -> str:
|
||||
"""Send one page image to DeepInfra olmOCR and return the extracted text."""
|
||||
b64 = image_to_base64(image_path)
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": MODEL,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{b64}"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"Below is a scanned page from a Canadian Army WWII war diary or supplementary document, "
|
||||
"circa 1944-1945. Extract all text exactly as it appears, preserving layout, "
|
||||
"column structure, dates, grid references, unit names, and abbreviations. "
|
||||
"Output clean Markdown. Do not add commentary or summaries."
|
||||
),
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
"max_tokens": 4096,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
for attempt in range(1, MAX_RETRIES + 1):
|
||||
try:
|
||||
response = requests.post(DEEPINFRA_API_URL, headers=headers, json=payload, timeout=120)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data["choices"][0]["message"]["content"]
|
||||
except requests.exceptions.HTTPError as e:
|
||||
print(f" HTTP error on page {page_num}, attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return f"[OCR FAILED - page {page_num} - HTTP error: {e}]"
|
||||
except Exception as e:
|
||||
print(f" Error on page {page_num}, attempt {attempt}/{MAX_RETRIES}: {e}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
return f"[OCR FAILED - page {page_num} - Error: {e}]"
|
||||
|
||||
|
||||
def ocr_pdf(pdf_path: Path, output_dir: Path, api_key: str) -> Path:
|
||||
"""OCR a single PDF and write output to a Markdown file."""
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Processing: {pdf_path.name}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
output_md = output_dir / f"{pdf_path.stem}_olmocr.md"
|
||||
|
||||
# Skip if already done
|
||||
if output_md.exists():
|
||||
print(f" Already processed — skipping. Delete {output_md.name} to reprocess.")
|
||||
return output_md
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Render PDF pages to images
|
||||
print(f" Rendering PDF pages to images at {DPI} DPI...")
|
||||
try:
|
||||
images = pdf_to_images(pdf_path, tmp_path)
|
||||
except RuntimeError as e:
|
||||
print(f" ERROR: {e}")
|
||||
return output_md
|
||||
|
||||
total_pages = len(images)
|
||||
print(f" Found {total_pages} pages.")
|
||||
|
||||
all_text = [
|
||||
f"# {pdf_path.stem}\n",
|
||||
f"*OCR'd by olmOCR ({MODEL}) via DeepInfra*\n",
|
||||
f"*Source: {pdf_path.name} — {total_pages} pages*\n",
|
||||
"---\n",
|
||||
]
|
||||
|
||||
for i, image_path in enumerate(images, start=1):
|
||||
print(f" Page {i}/{total_pages}...", end=" ", flush=True)
|
||||
page_text = ocr_page(image_path, api_key, i)
|
||||
all_text.append(f"\n---\n## Page {i}\n\n{page_text}\n")
|
||||
print("done")
|
||||
|
||||
# Small delay to avoid hammering the API
|
||||
if i < total_pages:
|
||||
time.sleep(0.5)
|
||||
|
||||
# Write output
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
with open(output_md, "w", encoding="utf-8") as f:
|
||||
f.write("\n".join(all_text))
|
||||
|
||||
print(f" Saved: {output_md}")
|
||||
return output_md
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="OCR war diary PDFs using olmOCR via DeepInfra.")
|
||||
parser.add_argument("--input_dir", required=True, help="Folder containing PDF files to process.")
|
||||
parser.add_argument("--output_dir", required=True, help="Folder to write Markdown output files.")
|
||||
parser.add_argument("--api_key", required=True, help="Your DeepInfra API key.")
|
||||
parser.add_argument("--single", default=None, help="Process a single PDF file instead of a folder.")
|
||||
args = parser.parse_args()
|
||||
|
||||
input_dir = Path(args.input_dir)
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if args.single:
|
||||
pdfs = [Path(args.single)]
|
||||
else:
|
||||
pdfs = sorted(input_dir.glob("*.pdf"))
|
||||
|
||||
if not pdfs:
|
||||
print(f"No PDF files found in {input_dir}")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"Found {len(pdfs)} PDF(s) to process.")
|
||||
print(f"Output directory: {output_dir}")
|
||||
|
||||
results = []
|
||||
for pdf_path in pdfs:
|
||||
output_md = ocr_pdf(pdf_path, output_dir, args.api_key)
|
||||
results.append(output_md)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Complete. {len(results)} file(s) processed.")
|
||||
print(f"Output files:")
|
||||
for r in results:
|
||||
print(f" {r}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user