breakpilot-lehrer/klausur-service/backend/ocr/pipeline/orientation_api.py

"""
Orientation & Page-Split API endpoints (Steps 1 and 1b of OCR Pipeline).
"""

import logging
import time
from typing import Any, Dict

import cv2
from fastapi import APIRouter, HTTPException

from cv_vocab_pipeline import detect_and_fix_orientation
from .page_crop import detect_page_splits
from .session_store import update_session_db

from .orientation_crop_helpers import ensure_cached, append_pipeline_log
from .page_sub_sessions import create_page_sub_sessions_full

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/v1/ocr-pipeline", tags=["ocr-pipeline"])


# ---------------------------------------------------------------------------
# Step 1: Orientation
# ---------------------------------------------------------------------------

@router.post("/sessions/{session_id}/orientation")
async def detect_orientation(session_id: str):
    """Detect and fix 90/180/270 degree rotations from scanners.

    Reads the original image, applies orientation correction,
    stores the result as oriented_png.
    """
    cached = await ensure_cached(session_id)

    img_bgr = cached.get("original_bgr")
    if img_bgr is None:
        raise HTTPException(status_code=400, detail="Original image not available")

    t0 = time.time()

    # Detect and fix orientation
    oriented_bgr, orientation_deg = detect_and_fix_orientation(img_bgr.copy())

    duration = time.time() - t0

    orientation_result = {
        "orientation_degrees": orientation_deg,
        "corrected": orientation_deg != 0,
        "duration_seconds": round(duration, 2),
    }

    # Encode oriented image
    success, png_buf = cv2.imencode(".png", oriented_bgr)
    oriented_png = png_buf.tobytes() if success else b""

    # Update cache
    cached["oriented_bgr"] = oriented_bgr
    cached["orientation_result"] = orientation_result

    # Persist to DB
    await update_session_db(
        session_id,
        oriented_png=oriented_png,
        orientation_result=orientation_result,
        current_step=2,
    )

    logger.info(
        "OCR Pipeline: orientation session %s: %d° (%s) in %.2fs",
        session_id, orientation_deg,
        "corrected" if orientation_deg else "no change",
        duration,
    )

    await append_pipeline_log(session_id, "orientation", {
        "orientation_degrees": orientation_deg,
        "corrected": orientation_deg != 0,
    }, duration_ms=int(duration * 1000))

    h, w = oriented_bgr.shape[:2]
    return {
        "session_id": session_id,
        **orientation_result,
        "image_width": w,
        "image_height": h,
        "oriented_image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/oriented",
    }


# ---------------------------------------------------------------------------
# Step 1b: Page-split detection — runs AFTER orientation, BEFORE deskew
# ---------------------------------------------------------------------------

@router.post("/sessions/{session_id}/page-split")
async def detect_page_split(session_id: str):
    """Detect if the image is a double-page book spread and split into sub-sessions.

    Must be called **after orientation** (step 1) and **before deskew** (step 2).
    Each sub-session receives the raw page region and goes through the full
    pipeline (deskew -> dewarp -> crop -> columns -> rows -> words -> grid)
    independently, so each page gets its own deskew correction.

    Returns ``{"multi_page": false}`` if only one page is detected.
    """
    cached = await ensure_cached(session_id)

    # Use oriented (preferred), fall back to original
    img_bgr = next(
        (v for k in ("oriented_bgr", "original_bgr")
         if (v := cached.get(k)) is not None),
        None,
    )
    if img_bgr is None:
        raise HTTPException(status_code=400, detail="No image available for page-split detection")

    t0 = time.time()
    page_splits = detect_page_splits(img_bgr)
    used_original = False

    if not page_splits or len(page_splits) < 2:
        # Orientation may have rotated a landscape double-page spread to
        # portrait.  Try the original (pre-orientation) image as fallback.
        orig_bgr = cached.get("original_bgr")
        if orig_bgr is not None and orig_bgr is not img_bgr:
            page_splits_orig = detect_page_splits(orig_bgr)
            if page_splits_orig and len(page_splits_orig) >= 2:
                logger.info(
                    "OCR Pipeline: page-split session %s: spread detected on "
                    "ORIGINAL (orientation rotated it away)",
                    session_id,
                )
                img_bgr = orig_bgr
                page_splits = page_splits_orig
                used_original = True

    if not page_splits or len(page_splits) < 2:
        duration = time.time() - t0
        logger.info(
            "OCR Pipeline: page-split session %s: single page (%.2fs)",
            session_id, duration,
        )
        return {
            "session_id": session_id,
            "multi_page": False,
            "duration_seconds": round(duration, 2),
        }

    # Multi-page spread detected — create sub-sessions for full pipeline.
    # start_step=2 means "ready for deskew" (orientation already applied).
    # start_step=1 means "needs orientation too" (split from original image).
    start_step = 1 if used_original else 2
    sub_sessions = await create_page_sub_sessions_full(
        session_id, cached, img_bgr, page_splits, start_step=start_step,
    )
    duration = time.time() - t0

    split_info: Dict[str, Any] = {
        "multi_page": True,
        "page_count": len(page_splits),
        "page_splits": page_splits,
        "used_original": used_original,
        "duration_seconds": round(duration, 2),
    }

    # Mark parent session as split and hidden from session list
    await update_session_db(session_id, crop_result=split_info, status='split')
    cached["crop_result"] = split_info

    await append_pipeline_log(session_id, "page_split", {
        "multi_page": True,
        "page_count": len(page_splits),
    }, duration_ms=int(duration * 1000))

    logger.info(
        "OCR Pipeline: page-split session %s: %d pages detected in %.2fs",
        session_id, len(page_splits), duration,
    )

    h, w = img_bgr.shape[:2]
    return {
        "session_id": session_id,
        **split_info,
        "image_width": w,
        "image_height": h,
        "sub_sessions": sub_sessions,
    }