""" 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 ocr_pipeline_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, }