304 lines
10 KiB
Python
304 lines
10 KiB
Python
"""
|
|
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() |