Create 02_profile_pdfs.py

smoke_signal/scripts/02_profile_pdfs.py

@@ -0,0 +1,250 @@
+#!/usr/bin/env python3
+"""
+Smoke Signal — Stage 3: PDF Profiler & Renderer
+=================================================
+For every PDF in the source manifest (status=pending or profiled),
+this script:
+  1. Opens each PDF and inspects every page
+  2. Detects whether a page has embedded text, is image-only, or hybrid
+  3. Detects rotation, skew hints, two-page spreads, dimensions
+  4. Routes each page: embedded_text | ocr | hybrid
+  5. Renders OCR-candidate pages to PNG at 300 DPI (450 DPI fallback)
+  6. Saves page_profile.json per book
+  7. Updates source_manifest.csv status to 'rendered'
+
+Usage:
+    python scripts/02_profile_pdfs.py
+    python scripts/02_profile_pdfs.py --book-id SS-BOOK-0001
+    python scripts/02_profile_pdfs.py --batch-id SS-BATCH-001
+    python scripts/02_profile_pdfs.py --dry-run
+"""
+
+import argparse
+import csv
+import json
+import sys
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+ROOT         = Path(__file__).resolve().parents[1]
+SOURCE_DIR   = ROOT / "source_pdfs"
+MANIFEST_CSV = ROOT / "manifest" / "source_manifest.csv"
+RENDERS_DIR  = ROOT / "renders"
+PROFILES_DIR = ROOT / "manifest" / "page_profiles"
+LOGS_DIR     = ROOT / "logs"
+
+PROFILES_DIR.mkdir(parents=True, exist_ok=True)
+RENDERS_DIR.mkdir(parents=True, exist_ok=True)
+LOGS_DIR.mkdir(parents=True, exist_ok=True)
+
+CONFIG = {
+    "config_version":        "ss_profiler_v0.1",
+    "render_dpi_baseline":   300,
+    "render_dpi_fallback":   450,
+    "text_char_threshold":   20,
+    "spread_aspect_ratio":   1.6,
+    "image_coverage_threshold": 0.15,
+    "skip_statuses":         ["exported", "quarantined"],
+    "eligible_statuses":     ["pending", "profiled"],
+}
+
+
+def load_manifest() -> dict:
+    records = {}
+    if not MANIFEST_CSV.exists():
+        return records
+    with open(MANIFEST_CSV, newline="", encoding="utf-8") as f:
+        for row in csv.DictReader(f):
+            if row.get("book_id"):
+                records[row["book_id"]] = row
+    return records
+
+
+def save_manifest(records: dict) -> None:
+    fields = [
+        "book_id", "source_id", "filename", "sha256", "file_size_bytes",
+        "page_count", "rights_class", "source_location", "acquisition_date",
+        "status", "allowed_use", "notes"
+    ]
+    rows = sorted(records.values(), key=lambda r: r.get("book_id", ""))
+    with open(MANIFEST_CSV, "w", newline="", encoding="utf-8") as f:
+        writer = csv.DictWriter(f, fieldnames=fields)
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+def _import_fitz():
+    try:
+        import fitz
+        return fitz
+    except ImportError:
+        print("  [error] PyMuPDF (fitz) not installed. Run: pip install pymupdf")
+        sys.exit(1)
+
+
+def analyse_page(page, fitz, config: dict) -> dict:
+    rect   = page.rect
+    width  = rect.width
+    height = rect.height
+    rotation = page.rotation
+    text = page.get_text("text").strip()
+    char_count = len(text)
+    has_embedded_text = char_count >= config["text_char_threshold"]
+    image_list = page.get_images(full=True)
+    has_images = len(image_list) > 0
+    image_coverage = 0.0
+    page_area = width * height
+    if has_images and page_area > 0:
+        for img in image_list:
+            try:
+                rects = page.get_image_rects(img[0])
+                for r in rects:
+                    image_coverage += abs(r.width * r.height) / page_area
+            except Exception:
+                pass
+        image_coverage = min(image_coverage, 1.0)
+    if has_embedded_text and not has_images:
+        route = "embedded_text"
+    elif has_embedded_text and has_images:
+        route = "hybrid" if image_coverage >= config["image_coverage_threshold"] else "embedded_text"
+    else:
+        route = "ocr"
+    aspect = width / height if height > 0 else 0
+    is_spread = aspect >= config["spread_aspect_ratio"]
+    warnings = []
+    if rotation not in (0, 360):
+        warnings.append(f"rotation_{rotation}deg")
+    if is_spread:
+        warnings.append("possible_two_page_spread")
+    if route == "ocr" and char_count == 0 and not has_images:
+        warnings.append("blank_or_undetectable_page")
+    return {
+        "width_pt": round(width, 2), "height_pt": round(height, 2),
+        "rotation_deg": rotation, "char_count": char_count,
+        "has_embedded_text": has_embedded_text, "has_images": has_images,
+        "image_count": len(image_list), "image_coverage": round(image_coverage, 3),
+        "is_spread": is_spread, "route": route, "warnings": warnings,
+    }
+
+
+def render_page(page, book_id: str, page_num: int, dpi: int, render_dir: Path) -> Optional[str]:
+    book_render_dir = render_dir / book_id
+    book_render_dir.mkdir(parents=True, exist_ok=True)
+    filename = f"{book_id}_page_{page_num:04d}_{dpi}dpi.png"
+    out_path = book_render_dir / filename
+    try:
+        mat = page.fitz_module.Matrix(dpi / 72, dpi / 72)
+        pix = page.get_pixmap(matrix=mat, alpha=False)
+        pix.save(str(out_path))
+        return str(out_path.relative_to(ROOT))
+    except Exception:
+        return None
+
+
+def profile_book(record: dict, dry_run: bool = False):
+    book_id  = record["book_id"]
+    filename = record["filename"]
+    pdf_path = SOURCE_DIR / filename
+    print(f"\n  [{book_id}] {filename}")
+    if not pdf_path.exists():
+        return record, None, {"error": "file_not_found"}
+    if record.get("rights_class") in ("unknown", "excluded"):
+        return record, None, {"error": "rights_blocked"}
+    fitz = _import_fitz()
+    try:
+        doc = fitz.open(str(pdf_path))
+    except Exception as e:
+        return record, None, {"error": str(e)}
+    page_count = doc.page_count
+    print(f"    Pages: {page_count}")
+    pages = []
+    render_errors = []
+    route_counts = {"embedded_text": 0, "ocr": 0, "hybrid": 0}
+    for i in range(page_count):
+        page_num = i + 1
+        page = doc[i]
+        page.fitz_module = fitz
+        profile = analyse_page(page, fitz, CONFIG)
+        profile["page_number"] = page_num
+        render_path = None
+        render_dpi  = None
+        if not dry_run and profile["route"] in ("ocr", "hybrid"):
+            dpi = CONFIG["render_dpi_baseline"]
+            render_path = render_page(page, book_id, page_num, dpi, RENDERS_DIR)
+            if render_path is None:
+                dpi = CONFIG["render_dpi_fallback"]
+                render_path = render_page(page, book_id, page_num, dpi, RENDERS_DIR)
+                if render_path is None:
+                    render_errors.append(page_num)
+                else:
+                    render_dpi = dpi
+            else:
+                render_dpi = dpi
+        profile["render_path"] = render_path
+        profile["render_dpi"]  = render_dpi
+        pages.append(profile)
+        route_counts[profile["route"]] += 1
+    doc.close()
+    book_profile = {
+        "book_id": book_id, "filename": filename,
+        "source_hash": record.get("sha256", ""),
+        "page_count": page_count, "config_version": CONFIG["config_version"],
+        "profiled_at": datetime.utcnow().isoformat() + "Z",
+        "route_summary": route_counts, "render_errors": render_errors, "pages": pages,
+    }
+    if not dry_run:
+        profile_path = PROFILES_DIR / f"{book_id}_page_profile.json"
+        with open(profile_path, "w", encoding="utf-8") as f:
+            json.dump(book_profile, f, indent=2)
+        print(f"    Profile saved -> {profile_path.relative_to(ROOT)}")
+    total_ocr = route_counts["ocr"] + route_counts["hybrid"]
+    print(f"    Routes: embedded={route_counts['embedded_text']} | ocr={route_counts['ocr']} | hybrid={route_counts['hybrid']}")
+    record["page_count"] = page_count
+    record["status"]     = "profiled" if not dry_run else record["status"]
+    return record, book_profile, {}
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Smoke Signal Stage 3: PDF Profiler")
+    parser.add_argument("--book-id", help="Profile a single book by ID")
+    parser.add_argument("--batch-id", help="Tag this run with a batch ID")
+    parser.add_argument("--dry-run", action="store_true")
+    parser.add_argument("--all", action="store_true")
+    args = parser.parse_args()
+    run_id  = args.batch_id or f"SS-RUN-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}"
+    dry_run = args.dry_run
+    print(f"\nSmoke Signal Stage 3 | {run_id} | config={CONFIG['config_version']}")
+    manifest = load_manifest()
+    if not manifest:
+        print("Manifest empty. Run 01_register_sources.py first.")
+        sys.exit(1)
+    eligible = CONFIG["eligible_statuses"] if not args.all else ["pending","profiled","rendered"]
+    books = [manifest[args.book_id]] if args.book_id else [r for r in manifest.values() if r.get("status") in eligible]
+    if not books:
+        print(f"No books with status in {eligible}.")
+        sys.exit(0)
+    print(f"Books to profile: {len(books)}")
+    results = []
+    t_start = time.time()
+    for record in books:
+        updated, profile, error = profile_book(record, dry_run=dry_run)
+        result = {"book_id": record["book_id"], "filename": record["filename"], "error": error}
+        if profile:
+            result.update({"route_summary": profile["route_summary"], "page_count": profile["page_count"]})
+            manifest[record["book_id"]] = updated
+        results.append(result)
+    if not dry_run:
+        save_manifest(manifest)
+        log_path = LOGS_DIR / f"{run_id}_profile_run.json"
+        with open(log_path, "w") as f:
+            json.dump({"run_id": run_id, "config": CONFIG, "results": results}, f, indent=2)
+    elapsed = round(time.time() - t_start, 1)
+    succeeded = sum(1 for r in results if not r.get("error"))
+    print(f"\nDone: {succeeded}/{len(results)} succeeded in {elapsed}s")
+    print("Next: run 03_ocr_bakeoff.py on calibration corpus")
+
+
+if __name__ == "__main__":
+    main()