smoke_signal/scripts/04_region_detector.py

#!/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()