Upload 04_region_detector.py

smoke_signal/scripts/04_region_detector.py

@@ -0,0 +1,452 @@
+#!/usr/bin/env python3
+"""
+Smoke Signal — Stage 5: Region Detection & Reading Order Resolver
+==================================================================
+Takes OCR raw output from Stage 4 and:
+  1. Classifies each detected region (narration, speech bubble, caption, etc.)
+  2. Assigns reading order using picture-book page logic
+  3. Detects two-page spreads and handles them correctly
+  4. Flags uncertain ordering rather than forcing false certainty
+  5. Saves ordered_regions.json per book
+  6. Creates visual overlay metadata for human review
+
+Region classes:
+    narration | dialogue-speech-bubble | title | subtitle | caption |
+    sign-label | page-number | copyright-legal | publisher-imprint | decorative-uncertain
+
+Usage:
+    python scripts/04_region_detector.py
+    python scripts/04_region_detector.py --book-id SS-BOOK-0001
+    python scripts/04_region_detector.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
+
+# ── Paths ──────────────────────────────────────────────────────────────────────
+ROOT         = Path(__file__).resolve().parents[1]
+MANIFEST_CSV = ROOT / "manifest" / "source_manifest.csv"
+PROFILES_DIR = ROOT / "manifest" / "page_profiles"
+OCR_RAW_DIR  = ROOT / "ocr_raw"
+REGIONS_DIR  = ROOT / "regions"
+LOGS_DIR     = ROOT / "logs"
+REVIEW_DIR   = ROOT / "review"
+
+REGIONS_DIR.mkdir(parents=True, exist_ok=True)
+LOGS_DIR.mkdir(parents=True, exist_ok=True)
+
+# ── Config ────────────────────────────────────────────────────────────────────
+CONFIG = {
+    "config_version":       "ss_regions_v0.1",
+    "page_number_max_chars": 4,     # regions with <= 4 chars near top/bottom = page number
+    "copyright_keywords":   ["©", "copyright", "all rights reserved", "isbn", "printed in"],
+    "publisher_keywords":   ["published by", "first published", "edition", "press", "publishers"],
+    "title_y_threshold":    0.20,   # top 20% of page = likely title zone
+    "footer_y_threshold":   0.85,   # bottom 15% of page = likely footer zone
+    "spread_gap_threshold": 0.45,   # x-center between 0.45-0.55 = possible spread gutter
+    "min_story_chars":      10,     # below this = not story text
+    "eligible_statuses":    ["ocred"],
+}
+
+REGION_CLASSES = [
+    "narration", "dialogue-speech-bubble", "title", "subtitle",
+    "caption", "sign-label", "page-number", "copyright-legal",
+    "publisher-imprint", "decorative-uncertain",
+]
+
+
+# ── Manifest I/O ───────────────────────────────────────────────────────────────
+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 load_page_profile(book_id: str) -> Optional[dict]:
+    path = PROFILES_DIR / f"{book_id}_page_profile.json"
+    return json.load(open(path)) if path.exists() else None
+
+
+def load_ocr_raw(book_id: str) -> Optional[dict]:
+    path = OCR_RAW_DIR / book_id / f"{book_id}_ocr_raw.json"
+    return json.load(open(path)) if path.exists() else None
+
+
+# ── Region classifier ─────────────────────────────────────────────────────────
+def classify_region(region: dict, page_height: float, page_width: float) -> str:
+    """
+    Heuristic region classification based on:
+    - Text content (keywords)
+    - Position on page (y-coordinate normalised 0-1)
+    - Text length
+    """
+    text  = region.get("text", "").strip()
+    lower = text.lower()
+    bbox  = region.get("bbox", [0, 0, 0, 0])  # [x0, y0, x1, y1]
+
+    if not text:
+        return "decorative-uncertain"
+
+    # Normalise bbox to 0-1 relative to page dimensions
+    if page_height > 0 and page_width > 0:
+        y_center  = ((bbox[1] + bbox[3]) / 2) / page_height
+        x_center  = ((bbox[0] + bbox[2]) / 2) / page_width
+    else:
+        y_center  = 0.5
+        x_center  = 0.5
+
+    char_count = len(text)
+
+    # ── Page number ──────────────────────────────────────────────────────────
+    if (char_count <= CONFIG["page_number_max_chars"] and text.strip().isdigit() and
+            (y_center < 0.12 or y_center > CONFIG["footer_y_threshold"])):
+        return "page-number"
+
+    # ── Copyright / legal ────────────────────────────────────────────────────
+    if any(kw in lower for kw in CONFIG["copyright_keywords"]):
+        return "copyright-legal"
+
+    # ── Publisher imprint ─────────────────────────────────────────────────────
+    if any(kw in lower for kw in CONFIG["publisher_keywords"]):
+        return "publisher-imprint"
+
+    # ── Title zone (top of page, short text) ─────────────────────────────────
+    if y_center < CONFIG["title_y_threshold"] and char_count < 60:
+        return "title"
+
+    # ── Speech bubble heuristic ───────────────────────────────────────────────
+    # Dialogue typically has quotes, said verbs, or is short and mid-page
+    has_quotes = any(c in text for c in ('"', '"', '"', "'", "«", "»"))
+    if has_quotes and char_count < 120:
+        return "dialogue-speech-bubble"
+
+    # ── Caption (short, bottom portion of page) ───────────────────────────────
+    if y_center > 0.75 and char_count < 80 and not text.isdigit():
+        return "caption"
+
+    # ── Sign / label (very short, anywhere) ──────────────────────────────────
+    if char_count < 20 and not text.isdigit():
+        return "sign-label"
+
+    # ── Decorative / uncertain (too short for story) ─────────────────────────
+    if char_count < CONFIG["min_story_chars"]:
+        return "decorative-uncertain"
+
+    # ── Default: narration ────────────────────────────────────────────────────
+    return "narration"
+
+
+# ── Reading order resolver ────────────────────────────────────────────────────
+def resolve_reading_order(regions: list, page_width: float, is_spread: bool) -> list:
+    """
+    Sort regions into picture-book reading order.
+
+    Rules:
+    - Standard page: top-to-bottom, left-to-right within same y-band
+    - Two-page spread: left page top-to-bottom, then right page top-to-bottom
+    - Page numbers and copyright always last
+    - Uncertain order is flagged, not forced
+    """
+    LAST_CLASSES = {"page-number", "copyright-legal", "publisher-imprint"}
+
+    story_regions  = [r for r in regions if r.get("region_class") not in LAST_CLASSES]
+    footer_regions = [r for r in regions if r.get("region_class") in LAST_CLASSES]
+
+    if is_spread and page_width > 0:
+        # Split into left and right page
+        mid_x = page_width / 2
+        left  = [r for r in story_regions if ((r["bbox"][0] + r["bbox"][2]) / 2) < mid_x]
+        right = [r for r in story_regions if ((r["bbox"][0] + r["bbox"][2]) / 2) >= mid_x]
+
+        def sort_key(r):
+            return (r["bbox"][1], r["bbox"][0])  # y then x
+
+        ordered = sorted(left, key=sort_key) + sorted(right, key=sort_key)
+    else:
+        # Standard: top-to-bottom with left-to-right tiebreak
+        y_band_size = 50  # points — regions within this y-band are on same "line"
+
+        def sort_key(r):
+            y_top = r["bbox"][1]
+            band  = round(y_top / y_band_size)
+            return (band, r["bbox"][0])  # band then x
+
+        ordered = sorted(story_regions, key=sort_key)
+
+    # Add reading order index
+    full_ordered = []
+    for i, region in enumerate(ordered + footer_regions):
+        region = dict(region)
+        region["reading_order"] = i + 1
+        region["reading_order_uncertain"] = False
+        full_ordered.append(region)
+
+    # Flag uncertainty for overlapping regions
+    for i in range(1, len(full_ordered)):
+        prev = full_ordered[i - 1]
+        curr = full_ordered[i]
+        # If bboxes overlap significantly in y, order is uncertain
+        prev_y_bottom = prev["bbox"][3]
+        curr_y_top    = curr["bbox"][1]
+        if curr_y_top < prev_y_bottom - 10:  # significant overlap
+            curr["reading_order_uncertain"] = True
+
+    return full_ordered
+
+
+# ── Per-page region processing ────────────────────────────────────────────────
+def process_page_regions(
+    book_id: str,
+    page_num: int,
+    ocr_page: dict,
+    page_profile: dict,
+) -> dict:
+    """Process regions for a single page."""
+
+    page_width  = page_profile.get("width_pt", 595)
+    page_height = page_profile.get("height_pt", 842)
+    is_spread   = page_profile.get("is_spread", False)
+    route       = page_profile.get("route", "ocr")
+
+    raw_regions = ocr_page.get("regions", [])
+
+    # Classify each region
+    classified = []
+    for region in raw_regions:
+        region = dict(region)
+        region["region_class"]  = classify_region(region, page_height, page_width)
+        region["region_id"]     = f"{book_id}_p{page_num:04d}_r{len(classified)+1:03d}"
+        region["source_book"]   = book_id
+        region["page_number"]   = page_num
+        classified.append(region)
+
+    # Resolve reading order
+    ordered = resolve_reading_order(classified, page_width, is_spread)
+
+    # Page-level stats
+    story_text = " ".join(
+        r["text"] for r in ordered
+        if r.get("region_class") in ("narration", "dialogue-speech-bubble", "caption")
+    )
+
+    return {
+        "book_id":          book_id,
+        "page_number":      page_num,
+        "route":            route,
+        "is_spread":        is_spread,
+        "page_width_pt":    page_width,
+        "page_height_pt":   page_height,
+        "region_count":     len(ordered),
+        "story_char_count": len(story_text),
+        "page_confidence":  ocr_page.get("page_confidence", 0),
+        "confidence_class": ocr_page.get("confidence_class", "review-required"),
+        "regions":          ordered,
+        "warnings":         page_profile.get("warnings", []),
+        "config_version":   CONFIG["config_version"],
+        "processed_at":     datetime.utcnow().isoformat() + "Z",
+    }
+
+
+# ── Per-book region detection ─────────────────────────────────────────────────
+def detect_regions_for_book(record: dict, dry_run: bool = False) -> tuple:
+    book_id  = record["book_id"]
+    filename = record["filename"]
+
+    print(f"\n  [{book_id}] {filename}")
+
+    if record.get("rights_class") in ("unknown", "excluded"):
+        return record, None, {"error": "rights_blocked"}
+
+    page_profile_data = load_page_profile(book_id)
+    if page_profile_data is None:
+        print(f"    ✗ No page profile. Run 02_profile_pdfs.py first.")
+        return record, None, {"error": "no_page_profile"}
+
+    ocr_data = load_ocr_raw(book_id)
+    if ocr_data is None:
+        print(f"    ✗ No OCR data. Run 03_ocr_bakeoff.py first.")
+        return record, None, {"error": "no_ocr_data"}
+
+    # Index page profiles by page number
+    page_profiles_by_num = {p["page_number"]: p for p in page_profile_data["pages"]}
+
+    # Index OCR pages by page number
+    ocr_pages_by_num = {p["page_number"]: p for p in ocr_data.get("pages", [])}
+
+    all_page_results = []
+    uncertain_pages  = []
+    empty_pages      = []
+
+    total_pages = page_profile_data["page_count"]
+
+    for page_num in range(1, total_pages + 1):
+        page_prof = page_profiles_by_num.get(page_num, {})
+        route     = page_prof.get("route", "embedded_text")
+
+        if route == "embedded_text":
+            # For embedded text pages, create minimal region record
+            all_page_results.append({
+                "book_id":        book_id,
+                "page_number":    page_num,
+                "route":          "embedded_text",
+                "region_count":   0,
+                "story_char_count": 0,
+                "regions":        [],
+                "note":           "embedded_text_page_regions_not_processed_here",
+            })
+            continue
+
+        ocr_page = ocr_pages_by_num.get(page_num)
+        if ocr_page is None:
+            empty_pages.append(page_num)
+            continue
+
+        if not dry_run:
+            page_result = process_page_regions(book_id, page_num, ocr_page, page_prof)
+        else:
+            page_result = {
+                "book_id": book_id, "page_number": page_num,
+                "route": route, "region_count": 0,
+                "regions": [], "note": "dry-run",
+            }
+
+        # Check for uncertain ordering
+        uncertain = [r for r in page_result.get("regions", []) if r.get("reading_order_uncertain")]
+        if uncertain:
+            uncertain_pages.append(page_num)
+
+        if page_result.get("story_char_count", 0) == 0 and route != "embedded_text":
+            empty_pages.append(page_num)
+
+        all_page_results.append(page_result)
+
+        region_count = page_result.get("region_count", 0)
+        conf         = page_result.get("page_confidence", 0)
+        print(f"    Page {page_num:3d}: {region_count} regions | conf={conf:.2f} | route={route}")
+
+    # ── Save ordered_regions.json ─────────────────────────────────────────────
+    book_regions = {
+        "book_id":        book_id,
+        "filename":       filename,
+        "source_hash":    record.get("sha256", ""),
+        "page_count":     total_pages,
+        "config_version": CONFIG["config_version"],
+        "processed_at":   datetime.utcnow().isoformat() + "Z",
+        "uncertain_pages": uncertain_pages,
+        "empty_story_pages": empty_pages,
+        "pages":          all_page_results,
+    }
+
+    if not dry_run:
+        regions_path = REGIONS_DIR / f"{book_id}_ordered_regions.json"
+        with open(regions_path, "w", encoding="utf-8") as f:
+            json.dump(book_regions, f, indent=2)
+        print(f"    Regions saved → {regions_path.relative_to(ROOT)}")
+
+        if uncertain_pages:
+            print(f"    ⚠️  Uncertain reading order on pages: {uncertain_pages}")
+        if empty_pages:
+            print(f"    ⚠️  Empty story pages: {empty_pages}")
+
+    record["status"] = "ocred"  # keep at ocred — reviewed/exported set later
+    return record, book_regions, {}
+
+
+# ── Main ───────────────────────────────────────────────────────────────────────
+def main():
+    parser = argparse.ArgumentParser(description="Smoke Signal — Stage 5: Region Detector")
+    parser.add_argument("--book-id",  help="Process 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")
+    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"\n{'='*60}")
+    print(f"  Smoke Signal — Stage 5: Region Detector")
+    print(f"  Run ID : {run_id}")
+    print(f"  Config : {CONFIG['config_version']}")
+    if dry_run:
+        print(f"  Mode   : DRY RUN")
+    print(f"{'='*60}")
+
+    manifest = load_manifest()
+    if not manifest:
+        print("\n  Manifest empty. Run earlier stages first.")
+        sys.exit(1)
+
+    books = (
+        [manifest[args.book_id]] if args.book_id and args.book_id in manifest
+        else [r for r in manifest.values() if r.get("status") in CONFIG["eligible_statuses"]]
+    )
+
+    if not books:
+        print(f"\n  No books with status in {CONFIG['eligible_statuses']}.")
+        print("  Run 03_ocr_bakeoff.py first.")
+        sys.exit(0)
+
+    print(f"\n  Books to process: {len(books)}")
+
+    results = []
+    t_start = time.time()
+
+    for record in books:
+        updated, book_regions, error = detect_regions_for_book(record, dry_run=dry_run)
+
+        result = {
+            "book_id":  record["book_id"],
+            "filename": record["filename"],
+            "error":    error,
+        }
+        if book_regions:
+            result["uncertain_pages"]    = book_regions.get("uncertain_pages", [])
+            result["empty_story_pages"]  = book_regions.get("empty_story_pages", [])
+            manifest[record["book_id"]] = updated
+
+        results.append(result)
+
+    if not dry_run:
+        save_manifest(manifest)
+        log_path = LOGS_DIR / f"{run_id}_regions_run.json"
+        with open(log_path, "w") as f:
+            json.dump({"run_id": run_id, "config": CONFIG, "results": results}, f, indent=2)
+        print(f"\n  Run log → {log_path.relative_to(ROOT)}")
+
+    elapsed   = round(time.time() - t_start, 1)
+    succeeded = sum(1 for r in results if not r.get("error"))
+
+    print(f"\n{'─'*60}")
+    print(f"  Books processed : {len(results)}")
+    print(f"  Succeeded       : {succeeded}")
+    print(f"  Time            : {elapsed}s")
+    print(f"{'─'*60}")
+    print(f"\n  Next: Run 05_llm_normalise.py (Stage 7)\n")
+
+
+if __name__ == "__main__":
+    main()