server.py

"""
Unified Local-First Multi-Model OCR Benchmark Backend Server (FastAPI).
Provides endpoints for single model OCR, sequential local 7-model benchmarking,
bounding box region visualizations, word/character diffing, and consensus evaluation.
"""

import os
import sys
import time
import json
import logging
from typing import Optional, List, Dict, Any
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from PIL import Image
import io

# Ensure project root in sys.path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from adapters import (
    AVAILABLE_MODELS,
    MODEL_CATALOG,
    get_adapter_by_name,
    get_model_info,
)
from core.models import Region, RegionType, ALL_REGION_TYPES, REGION_COLORS
from core.visualizer import render_annotated_image, save_visualization, image_to_base64_jpeg
from core.region_classifier import compute_region_summary
from core.benchmark_runner import run_sequential_local_benchmark
from utils.pdf_utils import load_input_image_or_pdf
from utils.diff_engine import compute_diff, generate_consensus_matrix, compute_similarity_ratio
from utils.judge_engine import evaluate_ocr_benchmark_judge

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
logger = logging.getLogger("OCRBackendServer")

app = FastAPI(
    title="Local-First Multi-Model OCR & Region Visualizer API",
    description="Local pipeline running 7 OCR models with bounding box extraction, classification, and visual rendering.",
    version="3.1.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Ensure visualization directory exists and mount static files
VIS_DIR = os.path.abspath("outputs/visualizations")
os.makedirs(VIS_DIR, exist_ok=True)
app.mount("/outputs/visualizations", StaticFiles(directory=VIS_DIR), name="visualizations")


@app.get("/api/health")
def health_check():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "service": "Local-First Multi-Model OCR Backend",
        "models_available": AVAILABLE_MODELS,
        "region_types": ALL_REGION_TYPES,
        "region_colors": REGION_COLORS,
        "timestamp": time.time()
    }


@app.get("/api/models")
def get_models():
    """Returns the metadata catalog for all 7 OCR models."""
    return {
        "models": AVAILABLE_MODELS,
        "catalog": MODEL_CATALOG,
        "region_types": ALL_REGION_TYPES,
        "region_colors": REGION_COLORS
    }


class DiffRequest(BaseModel):
    text_a: str
    text_b: str
    mode: str = "word"  # 'word' or 'char'


@app.post("/api/diff")
def calculate_diff_endpoint(req: DiffRequest):
    """Calculates word-level or character-level diff between two texts."""
    diff_chunks = compute_diff(req.text_a, req.text_b, mode=req.mode)
    similarity = compute_similarity_ratio(req.text_a, req.text_b)
    sim_pct = round(similarity * 100, 1) if similarity is not None else None
    return {
        "similarity_pct": sim_pct,
        "mode": req.mode,
        "chunks": diff_chunks
    }


@app.post("/api/ocr")
async def run_single_ocr(
    file: UploadFile = File(...),
    model_name: str = Form("PP-OCRv5"),
    page: int = Form(1),
    prompt: Optional[str] = Form(None)
):
    """Executes OCR with a single model and returns text + bounding box regions."""
    if model_name not in AVAILABLE_MODELS:
        raise HTTPException(status_code=400, detail=f"Invalid model '{model_name}'. Available: {AVAILABLE_MODELS}")

    contents = await file.read()
    try:
        target_img, total_pages, _ = load_input_image_or_pdf(contents, page_index=max(0, page - 1))
    except Exception as e:
        logger.error(f"Error loading image/PDF: {e}")
        raise HTTPException(status_code=400, detail=f"Could not decode image or PDF page: {str(e)}")

    if target_img is None:
        raise HTTPException(status_code=400, detail="Could not decode image or PDF page.")

    adapter = get_adapter_by_name(model_name)
    kwargs = {}
    if prompt:
        kwargs["prompt"] = prompt

    result = adapter.process(target_img, **kwargs)

    # Render visualization if successful
    annotated_base64 = None
    if result.get("status") == "SUCCESS" and result.get("regions"):
        try:
            annotated_img = render_annotated_image(target_img, result["regions"])
            annotated_base64 = image_to_base64_jpeg(annotated_img)
        except Exception as e:
            logger.warning(f"Failed to render single OCR visualization: {e}")

    result["annotated_image_base64"] = annotated_base64
    return result


@app.post("/api/benchmark")
async def run_benchmark_all_models(
    file: UploadFile = File(...),
    page: int = Form(1)
):
    """
    Main Benchmark Workflow:
    1. Reads uploaded image/PDF in-memory (no file locking).
    2. Runs all 7 models sequentially on local hardware.
    3. Measures accurate inference time with time.perf_counter().
    4. Extracts bounding boxes and classifies regions across 10 categories.
    5. Generates annotated visualization images per model.
    6. Computes Word/Character Diff and Consensus Matrix on SUCCESSFUL models.
    """
    contents = await file.read()
    filename = file.filename or "document.png"

    try:
        target_img, total_pages, _ = load_input_image_or_pdf(contents, page_index=max(0, page - 1))
    except Exception as e:
        logger.error(f"Failed to decode uploaded document '{filename}': {e}")
        raise HTTPException(status_code=400, detail=f"Could not decode document '{filename}': {str(e)}")

    if target_img is None:
        raise HTTPException(status_code=400, detail="Could not read uploaded document image or PDF.")

    logger.info(f"Running 7-model benchmark on '{filename}' (Page {page}/{total_pages}, Size: {target_img.size})")

    benchmark_data = run_sequential_local_benchmark(
        image=target_img,
        output_dir=VIS_DIR,
        document_name=filename
    )

    benchmark_data["page_number"] = page
    benchmark_data["total_pages"] = total_pages

    # Run LLM-as-a-Judge cross-evaluation on successful models
    judge_data = evaluate_ocr_benchmark_judge(benchmark_data["results"], document_name=filename)
    benchmark_data["judge_verdict"] = judge_data

    return benchmark_data


if os.path.exists("index.html"):
    @app.get("/")
    def serve_frontend():
        return FileResponse("index.html")

if __name__ == "__main__":
    import uvicorn
    import argparse

    parser = argparse.ArgumentParser(description="OCR Space Server")
    parser.add_argument("--host", type=str, default=os.getenv("HOST", "0.0.0.0"), help="Host IP")
    parser.add_argument("--port", type=int, default=int(os.getenv("PORT", "7860")), help="Port number")
    args, _ = parser.parse_known_args()

    logger.info(f"Starting server on {args.host}:{args.port}")
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")