breakpilot-lehrer/klausur-service/backend/crop_api.py

"""
Crop API endpoints (Step 4 / UI index 3 of OCR Pipeline).

Auto-crop, manual crop, and skip-crop for scanner/book borders.
"""

import logging
import time
from typing import Any, Dict

import cv2
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

from page_crop import detect_and_crop_page, detect_page_splits
from ocr_pipeline_session_store import get_sub_sessions, update_session_db

from orientation_crop_helpers import ensure_cached, append_pipeline_log
from page_sub_sessions import create_page_sub_sessions

logger = logging.getLogger(__name__)

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


# ---------------------------------------------------------------------------
# Step 4 (UI index 3): Crop — runs after deskew + dewarp
# ---------------------------------------------------------------------------

@router.post("/sessions/{session_id}/crop")
async def auto_crop(session_id: str):
    """Auto-detect and crop scanner/book borders.

    Reads the dewarped image (post-deskew + dewarp, so the page is straight).
    Falls back to oriented -> original if earlier steps were skipped.

    If the image is a multi-page spread (e.g. book on scanner), it will
    automatically split into separate sub-sessions per page, crop each
    individually, and return the split info.
    """
    cached = await ensure_cached(session_id)

    # Use dewarped (preferred), fall back to oriented, then original
    img_bgr = next(
        (v for k in ("dewarped_bgr", "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 cropping")

    t0 = time.time()

    # --- Check for existing sub-sessions (from page-split step) ---
    # If page-split already created sub-sessions, skip multi-page detection
    # in the crop step. Each sub-session runs its own crop independently.
    existing_subs = await get_sub_sessions(session_id)
    if existing_subs:
        crop_result = cached.get("crop_result") or {}
        if crop_result.get("multi_page"):
            # Already split -- just return the existing info
            duration = time.time() - t0
            h, w = img_bgr.shape[:2]
            return {
                "session_id": session_id,
                **crop_result,
                "image_width": w,
                "image_height": h,
                "sub_sessions": [
                    {"id": s["id"], "name": s.get("name"), "page_index": s.get("box_index", i)}
                    for i, s in enumerate(existing_subs)
                ],
                "note": "Page split was already performed; each sub-session runs its own crop.",
            }

    # --- Multi-page detection (fallback for sessions that skipped page-split) ---
    page_splits = detect_page_splits(img_bgr)

    if page_splits and len(page_splits) >= 2:
        # Multi-page spread detected -- create sub-sessions
        sub_sessions = await create_page_sub_sessions(
            session_id, cached, img_bgr, page_splits,
        )
        duration = time.time() - t0

        crop_info: Dict[str, Any] = {
            "crop_applied": True,
            "multi_page": True,
            "page_count": len(page_splits),
            "page_splits": page_splits,
            "duration_seconds": round(duration, 2),
        }
        cached["crop_result"] = crop_info

        # Store the first page as the main cropped image for backward compat
        first_page = page_splits[0]
        first_bgr = img_bgr[
            first_page["y"]:first_page["y"] + first_page["height"],
            first_page["x"]:first_page["x"] + first_page["width"],
        ].copy()
        first_cropped, _ = detect_and_crop_page(first_bgr)
        cached["cropped_bgr"] = first_cropped

        ok, png_buf = cv2.imencode(".png", first_cropped)
        await update_session_db(
            session_id,
            cropped_png=png_buf.tobytes() if ok else b"",
            crop_result=crop_info,
            current_step=5,
            status='split',
        )

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

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

        h, w = first_cropped.shape[:2]
        return {
            "session_id": session_id,
            **crop_info,
            "image_width": w,
            "image_height": h,
            "cropped_image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/cropped",
            "sub_sessions": sub_sessions,
        }

    # --- Single page (normal) ---
    cropped_bgr, crop_info = detect_and_crop_page(img_bgr)

    duration = time.time() - t0
    crop_info["duration_seconds"] = round(duration, 2)
    crop_info["multi_page"] = False

    # Encode cropped image
    success, png_buf = cv2.imencode(".png", cropped_bgr)
    cropped_png = png_buf.tobytes() if success else b""

    # Update cache
    cached["cropped_bgr"] = cropped_bgr
    cached["crop_result"] = crop_info

    # Persist to DB
    await update_session_db(
        session_id,
        cropped_png=cropped_png,
        crop_result=crop_info,
        current_step=5,
    )

    logger.info(
        "OCR Pipeline: crop session %s: applied=%s format=%s in %.2fs",
        session_id, crop_info["crop_applied"],
        crop_info.get("detected_format", "?"),
        duration,
    )

    await append_pipeline_log(session_id, "crop", {
        "crop_applied": crop_info["crop_applied"],
        "detected_format": crop_info.get("detected_format"),
        "format_confidence": crop_info.get("format_confidence"),
    }, duration_ms=int(duration * 1000))

    h, w = cropped_bgr.shape[:2]
    return {
        "session_id": session_id,
        **crop_info,
        "image_width": w,
        "image_height": h,
        "cropped_image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/cropped",
    }


class ManualCropRequest(BaseModel):
    x: float       # percentage 0-100
    y: float       # percentage 0-100
    width: float   # percentage 0-100
    height: float  # percentage 0-100


@router.post("/sessions/{session_id}/crop/manual")
async def manual_crop(session_id: str, req: ManualCropRequest):
    """Manually crop using percentage coordinates."""
    cached = await ensure_cached(session_id)

    img_bgr = next(
        (v for k in ("dewarped_bgr", "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 cropping")

    h, w = img_bgr.shape[:2]

    # Convert percentages to pixels
    px_x = int(w * req.x / 100.0)
    px_y = int(h * req.y / 100.0)
    px_w = int(w * req.width / 100.0)
    px_h = int(h * req.height / 100.0)

    # Clamp
    px_x = max(0, min(px_x, w - 1))
    px_y = max(0, min(px_y, h - 1))
    px_w = max(1, min(px_w, w - px_x))
    px_h = max(1, min(px_h, h - px_y))

    cropped_bgr = img_bgr[px_y:px_y + px_h, px_x:px_x + px_w].copy()

    success, png_buf = cv2.imencode(".png", cropped_bgr)
    cropped_png = png_buf.tobytes() if success else b""

    crop_result = {
        "crop_applied": True,
        "crop_rect": {"x": px_x, "y": px_y, "width": px_w, "height": px_h},
        "crop_rect_pct": {"x": round(req.x, 2), "y": round(req.y, 2),
                          "width": round(req.width, 2), "height": round(req.height, 2)},
        "original_size": {"width": w, "height": h},
        "cropped_size": {"width": px_w, "height": px_h},
        "method": "manual",
    }

    cached["cropped_bgr"] = cropped_bgr
    cached["crop_result"] = crop_result

    await update_session_db(
        session_id,
        cropped_png=cropped_png,
        crop_result=crop_result,
        current_step=5,
    )

    ch, cw = cropped_bgr.shape[:2]
    return {
        "session_id": session_id,
        **crop_result,
        "image_width": cw,
        "image_height": ch,
        "cropped_image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/cropped",
    }


@router.post("/sessions/{session_id}/crop/skip")
async def skip_crop(session_id: str):
    """Skip cropping -- use dewarped (or oriented/original) image as-is."""
    cached = await ensure_cached(session_id)

    img_bgr = next(
        (v for k in ("dewarped_bgr", "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")

    h, w = img_bgr.shape[:2]

    # Store the dewarped image as cropped (identity crop)
    success, png_buf = cv2.imencode(".png", img_bgr)
    cropped_png = png_buf.tobytes() if success else b""

    crop_result = {
        "crop_applied": False,
        "skipped": True,
        "original_size": {"width": w, "height": h},
        "cropped_size": {"width": w, "height": h},
    }

    cached["cropped_bgr"] = img_bgr
    cached["crop_result"] = crop_result

    await update_session_db(
        session_id,
        cropped_png=cropped_png,
        crop_result=crop_result,
        current_step=5,
    )

    return {
        "session_id": session_id,
        **crop_result,
        "image_width": w,
        "image_height": h,
        "cropped_image_url": f"/api/v1/ocr-pipeline/sessions/{session_id}/image/cropped",
    }