ocr_pages_to_text.py

#!/usr/bin/env python3
"""
Run Tesseract OCR on page images and save one .txt file per page.

Requires: tesseract-ocr installed (apt install tesseract-ocr),
          and optionally: pip install pytesseract Pillow

Usage:
  python ocr_pages_to_text.py [--pages-dir DIR] [--output-dir DIR] [--lang LANG]
"""

import argparse
import re
import subprocess
import sys
from pathlib import Path
from typing import List, Optional

try:
    import pytesseract
    from PIL import Image
    HAS_PYTESSERACT = True
except ImportError:
    HAS_PYTESSERACT = False

# Image extensions we look for (order matters: prefer lossless/high quality)
IMAGE_GLOB = ["*.png", "*.jpg", "*.jpeg", "*.tiff", "*.tif"]


def natural_sort_key(p: Path) -> tuple:
    """Sort by numeric suffix (page_0001, page_002, etc.)."""
    m = re.search(r"(\d+)\s*$", p.stem)
    return (int(m.group(1)),) if m else (0,)


def find_page_images(pages_dir: Path) -> List[Path]:
    images = []
    for ext in IMAGE_GLOB:
        images.extend(pages_dir.glob(ext))
    return sorted(set(images), key=natural_sort_key)


def ocr_with_subprocess(image_path: Path, out_path: Path, lang: str = "eng") -> bool:
    """Use tesseract CLI (works without pytesseract)."""
    try:
        subprocess.run(
            ["tesseract", str(image_path), str(out_path.with_suffix("")), "-l", lang, "txt"],
            check=True,
            capture_output=True,
            timeout=120,
        )
        return True
    except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired) as e:
        print(f"Error OCR {image_path.name}: {e}", file=sys.stderr)
        return False


def ocr_with_pytesseract(image_path: Path, lang: str = "eng") -> Optional[str]:
    """Use pytesseract (better if Pillow available)."""
    if not HAS_PYTESSERACT:
        return None
    try:
        img = Image.open(image_path)
        return pytesseract.image_to_string(img, lang=lang)
    except Exception as e:
        print(f"Error OCR {image_path.name}: {e}", file=sys.stderr)
        return None


def main():
    p = argparse.ArgumentParser(description="OCR page images to text files.")
    p.add_argument("--pages-dir", type=Path, default=Path("stand-on-zanzibar-pages"), help="Directory with page images")
    p.add_argument("--output-dir", "-o", type=Path, default=None, help="Directory for .txt files (default: pages-dir)")
    p.add_argument("--lang", default="eng", help="Tesseract language (default: eng)")
    p.add_argument("--use-cli", action="store_true", help="Use tesseract CLI instead of pytesseract")
    args = p.parse_args()

    pages_dir = args.pages_dir.resolve()
    if not pages_dir.is_dir():
        p.error(f"Pages directory not found: {pages_dir}")

    output_dir = (args.output_dir or pages_dir).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    images = find_page_images(pages_dir)
    if not images:
        print("No page images found.", file=sys.stderr)
        sys.exit(1)

    print(f"Found {len(images)} page images. OCR with lang={args.lang}")

    for i, img_path in enumerate(images, 1):
        # Output filename: same stem, .txt (e.g. stand-on-zanzibar-page_001.txt)
        out_txt = output_dir / f"{img_path.stem}.txt"
        if out_txt.exists():
            print(f"Skip (exists): {out_txt.name}")
            continue

        if args.use_cli or not HAS_PYTESSERACT:
            ok = ocr_with_subprocess(img_path, out_txt, args.lang)
            if ok:
                print(f"OCR: {img_path.name} -> {out_txt.name}")
        else:
            text = ocr_with_pytesseract(img_path, args.lang)
            if text is not None:
                out_txt.write_text(text.strip(), encoding="utf-8")
                print(f"OCR: {img_path.name} -> {out_txt.name}")


if __name__ == "__main__":
    main()