""" ocr_wardiaries.py ----------------- Converts war diary PDFs to plain text using the olmOCR model hosted on DeepInfra. Bypasses the olmOCR pipeline GPU check by calling the API directly. Output format: plain text only — no markdown, no HTML, no tables. Each diary entry is transcribed as prose with a blank line between entries. 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" python ocr_wardiaries.py --single "C:/path/to/diary.pdf" --output_dir "C:/path/to/output" --api_key "YOUR_KEY" """ import argparse import base64 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 # --------------------------------------------------------------------------- # OCR prompt # --------------------------------------------------------------------------- # Plain text only — no markdown, no HTML. # War diary pages are tabular (Place / Date / Hour / Summary). # We transcribe each row as prose to avoid HTML table output # which causes downstream parsing problems. SYSTEM_PROMPT = """\ You are transcribing a scanned WWII Canadian Army war diary page. Output plain text only. No markdown. No HTML. No tables. No formatting symbols. Rules: - Transcribe ALL text visible on the page. Every word, every field, every stamp, every notation. - Do not decide what is important. Do not skip anything. Transcribe everything. - Do not correct spelling, grammar, or abbreviations. Transcribe exactly as written. - If the page has a table with columns (Place / Date / Hour / Summary), transcribe each row \ as continuous prose in this format: [DATE] [PLACE]: [SUMMARY TEXT] Example: 4 Sep 44 France, NEUVILLE Les DIEPPE MR 2468 Sheet: The main topic for the morning... - Separate multiple diary entries with a single blank line. - Keep grid references exactly as written (e.g. MR 2468, GR 442891, MR 0675 Sheet). - Keep military abbreviations exactly as written (e.g. Bn, Bde, HQ, OR, C.O., R de Mais, RHC). - Weather notes go on their own line starting with: Weather: - Transcribe page headers, form titles, stamps, handwritten notes, signatures, \ and margin notations — all of it. - If a page is genuinely blank with no text at all, output only: [NO DIARY CONTENT] - Do not add commentary, explanations, confidence notes, or summaries. - Do not output page numbers or headings. - Do not use asterisks, hashes, pipes, underscores, or any other formatting characters.\ """ USER_PROMPT = ( "Transcribe every word on this war diary page exactly as written. " "Plain text only — include all text visible on the page, nothing omitted." ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def pdf_to_images(pdf_path: Path, output_dir: Path, dpi: int = DPI, first_page: int = None, last_page: int = None) -> 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", ] if first_page is not None: cmd += ["-f", str(first_page)] if last_page is not None: cmd += ["-l", str(last_page)] cmd += [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}") 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 plain text.""" b64 = image_to_base64(image_path) headers = { "Content-Type": "application/json", "Authorization": f"Bearer {api_key}", } payload = { "model": MODEL, "messages": [ { "role": "system", "content": SYSTEM_PROMPT, }, { "role": "user", "content": [ { "type": "image_url", "image_url": { "url": f"data:image/png;base64,{b64}" }, }, { "type": "text", "text": USER_PROMPT, }, ], } ], "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, max_pages: int = None) -> Path: """OCR a single PDF and write output to a plain text 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.") print(f" 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, last_page=max_pages) except RuntimeError as e: print(f" ERROR: {e}") return output_md total_pages = len(images) print(f" Found {total_pages} pages to process.") # File header — plain text, no markdown all_text = [ f"# {pdf_path.stem}", f"OCR by olmOCR ({MODEL}) via DeepInfra", f"Source: {pdf_path.name} — {len(images)} page(s) processed", "", ] 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) # Each page is separated by a ## Page N marker so the viewer # can split pages correctly. The content itself is plain text. all_text.append(f"## Page {i}") all_text.append("") all_text.append(page_text.strip()) all_text.append("") 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 to plain text using olmOCR via DeepInfra." ) parser.add_argument( "--input_dir", default=None, help="Folder containing PDF files to process.", ) parser.add_argument( "--output_dir", required=True, help="Folder to write 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 whole folder.", ) parser.add_argument( "--max_pages", type=int, default=None, help="Only process the first N pages of each PDF (useful for testing).", ) args = parser.parse_args() output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) if args.single: pdfs = [Path(args.single)] elif args.input_dir: pdfs = sorted(Path(args.input_dir).glob("*.pdf")) else: print("Error: provide --input_dir or --single") sys.exit(1) if not pdfs: print(f"No PDF files found.") sys.exit(1) print(f"Found {len(pdfs)} PDF(s) to process.") print(f"Output directory: {output_dir}") print(f"Model: {MODEL}") print(f"DPI: {DPI}") results = [] for pdf_path in pdfs: output_md = ocr_pdf(pdf_path, output_dir, args.api_key, max_pages=args.max_pages) results.append(output_md) print(f"\n{'='*60}") print(f"Complete. {len(results)} file(s) processed.") print("Output files:") for r in results: print(f" {r}") if __name__ == "__main__": main()