nathan/AI-Prototype
1
0
Fork 0
Code Pull Requests Packages Projects Activity

early version of OCR viewer now integrated

Browse Source
This commit is contained in:
nathan
2026-05-13 10:03:38 -04:00
parent 77323821cb
commit c663627594
17 changed files with 5734 additions and 1517 deletions
Download Patch File Download Diff File Expand all files Collapse all files

162
scripts/ocr_wardiaries.py
View File

@@ -1,9 +1,12 @@
"""
ocr_wardiaries.py
-----------------
Converts war diary PDFs to Markdown using the olmOCR model hosted on DeepInfra.
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
@@ -11,11 +14,11 @@ 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 json
import os
import subprocess
import sys
@@ -31,31 +34,71 @@ from PIL import Image
# ---------------------------------------------------------------------------
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
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) -> list[Path]:
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",
str(pdf_path),
str(prefix),
]
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}")
# 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
@@ -67,7 +110,7 @@ def image_to_base64(image_path: Path) -> str:
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."""
"""Send one page image to DeepInfra olmOCR and return the extracted plain text."""
b64 = image_to_base64(image_path)
headers = {
@@ -78,6 +121,10 @@ def ocr_page(image_path: Path, api_key: str, page_num: int) -> str:
payload = {
"model": MODEL,
"messages": [
{
"role": "system",
"content": SYSTEM_PROMPT,
},
{
"role": "user",
"content": [
@@ -89,12 +136,7 @@ def ocr_page(image_path: Path, api_key: str, page_num: int) -> str:
},
{
"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."
),
"text": USER_PROMPT,
},
],
}
@@ -105,7 +147,12 @@ def ocr_page(image_path: Path, api_key: str, page_num: int) -> str:
for attempt in range(1, MAX_RETRIES + 1):
try:
response = requests.post(DEEPINFRA_API_URL, headers=headers, json=payload, timeout=120)
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"]
@@ -123,8 +170,8 @@ def ocr_page(image_path: Path, api_key: str, page_num: int) -> str:
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."""
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}")
@@ -133,7 +180,8 @@ def ocr_pdf(pdf_path: Path, output_dir: Path, api_key: str) -> Path:
# Skip if already done
if output_md.exists():
print(f" Already processed — skipping. Delete {output_md.name} to reprocess.")
print(f" Already processed — skipping.")
print(f" Delete {output_md.name} to reprocess.")
return output_md
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -142,25 +190,33 @@ def ocr_pdf(pdf_path: Path, output_dir: Path, api_key: str) -> Path:
# Render PDF pages to images
print(f" Rendering PDF pages to images at {DPI} DPI...")
try:
images = pdf_to_images(pdf_path, tmp_path)
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.")
print(f" Found {total_pages} pages to process.")
# File header — plain text, no markdown
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",
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)
all_text.append(f"\n---\n## Page {i}\n\n{page_text}\n")
# 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
@@ -181,40 +237,68 @@ def ocr_pdf(pdf_path: Path, output_dir: Path, api_key: str) -> Path:
# ---------------------------------------------------------------------------
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.")
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()
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)]
elif args.input_dir:
pdfs = sorted(Path(args.input_dir).glob("*.pdf"))
else:
pdfs = sorted(input_dir.glob("*.pdf"))
print("Error: provide --input_dir or --single")
sys.exit(1)
if not pdfs:
print(f"No PDF files found in {input_dir}")
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)
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(f"Output files:")
print("Output files:")
for r in results:
print(f" {r}")
if __name__ == "__main__":
main()
main()