""" 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()