breakpilot-core/paddleocr-service/main.py

"""PaddleOCR Remote Service — PP-OCRv5 Latin auf x86_64."""

import io
import logging
import os
import threading

# Disable oneDNN/MKLDNN before importing paddle — avoids
# ConvertPirAttribute2RuntimeAttribute errors on PaddlePaddle 3.x
os.environ["FLAGS_use_mkldnn"] = "0"
os.environ["FLAGS_use_onednn"] = "0"
os.environ["PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK"] = "1"

# Must set paddle flags before import
try:
    import paddle
    paddle.set_flags({"FLAGS_use_mkldnn": False})
except Exception:
    pass

import numpy as np
from fastapi import FastAPI, File, Header, HTTPException, UploadFile
from PIL import Image

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="PaddleOCR Service")

_engine = None
_ready = False
_loading = False
API_KEY = os.environ.get("PADDLEOCR_API_KEY", "")


def _load_model():
    """Load PaddleOCR model in background thread."""
    global _engine, _ready
    try:
        logger.info("Importing paddleocr...")
        from paddleocr import PaddleOCR

        logger.info("Import done. Loading PaddleOCR model...")
        # Try multiple init strategies for different PaddleOCR versions
        inits = [
            # PaddleOCR 3.x — disable MKLDNN via enable_mkldnn=False
            dict(lang="en", ocr_version="PP-OCRv5", use_textline_orientation=True, enable_mkldnn=False),
            # PaddleOCR 3.x without enable_mkldnn param
            dict(lang="en", ocr_version="PP-OCRv5", use_textline_orientation=True),
            # PaddleOCR 3.x with deprecated param
            dict(lang="en", ocr_version="PP-OCRv5", use_angle_cls=True),
            # PaddleOCR 2.8+ (en, no version)
            dict(lang="en", use_angle_cls=True, show_log=False),
        ]
        for i, kwargs in enumerate(inits):
            try:
                _engine = PaddleOCR(**kwargs)
                logger.info(f"PaddleOCR init succeeded with strategy {i}: {kwargs}")
                break
            except Exception as e:
                logger.info(f"PaddleOCR init strategy {i} failed: {e}")
        else:
            raise RuntimeError("All PaddleOCR init strategies failed")
        _ready = True
        logger.info("PaddleOCR model loaded successfully — ready to serve")
    except Exception as e:
        logger.error(f"Failed to load PaddleOCR model: {e}")


@app.on_event("startup")
def startup_load_model():
    """Start model loading in background so health check passes immediately."""
    global _loading
    _loading = True
    thread = threading.Thread(target=_load_model, daemon=True)
    thread.start()
    logger.info("Model loading started in background thread")


@app.get("/health")
def health():
    if _ready:
        return {"status": "ok", "model": "PP-OCRv5-latin"}
    if _loading:
        return {"status": "loading"}
    return {"status": "error"}


@app.post("/ocr")
async def ocr(
    file: UploadFile = File(...),
    x_api_key: str = Header(default=""),
):
    if API_KEY and x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")

    if not _ready:
        raise HTTPException(status_code=503, detail="Model still loading")

    img_bytes = await file.read()
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    img_np = np.array(img)

    try:
        result = _engine.ocr(img_np)
    except Exception as e:
        logger.error(f"OCR failed: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"OCR processing failed: {e}")

    if not result:
        return {"words": [], "image_width": img_np.shape[1], "image_height": img_np.shape[0]}

    # PaddleOCR 2.x returns: [[line, ...]] where line = [box, (text, conf)]
    # PaddleOCR 3.x returns: [{'text': ..., 'boxes': [...], 'rec_scores': ...}] or similar
    words = []
    try:
        lines = result[0] if isinstance(result, list) and result else result
        if not lines:
            return {"words": [], "image_width": img_np.shape[1], "image_height": img_np.shape[0]}

        for line in lines:
            if isinstance(line, dict):
                # PaddleOCR 3.x dict format
                text = str(line.get("text", line.get("rec_text", ""))).strip()
                conf = float(line.get("score", line.get("rec_score", 0)))
                box = line.get("boxes", line.get("dt_polys", []))
                if not text or not box:
                    continue
                # box might be [[x1,y1],[x2,y2],[x3,y3],[x4,y4]] or flat
                if isinstance(box[0], (list, tuple)):
                    x_min = min(p[0] for p in box)
                    y_min = min(p[1] for p in box)
                    x_max = max(p[0] for p in box)
                    y_max = max(p[1] for p in box)
                else:
                    x_min, y_min, x_max, y_max = box[0], box[1], box[2], box[3]
                words.append({
                    "text": text,
                    "left": int(x_min), "top": int(y_min),
                    "width": int(x_max - x_min), "height": int(y_max - y_min),
                    "conf": round(conf * 100 if conf <= 1 else conf, 1),
                })
            elif isinstance(line, (list, tuple)) and len(line) == 2:
                # PaddleOCR 2.x format: [box, (text, conf)]
                box, (text, conf) = line[0], line[1]
                x_min = min(p[0] for p in box)
                y_min = min(p[1] for p in box)
                x_max = max(p[0] for p in box)
                y_max = max(p[1] for p in box)
                words.append({
                    "text": str(text).strip(),
                    "left": int(x_min), "top": int(y_min),
                    "width": int(x_max - x_min), "height": int(y_max - y_min),
                    "conf": round(float(conf) * 100 if conf <= 1 else float(conf), 1),
                })
    except Exception as e:
        logger.error(f"Failed to parse OCR result: {e}. Raw: {str(result)[:500]}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"OCR result parsing failed: {e}")

    return {
        "words": words,
        "image_width": img_np.shape[1],
        "image_height": img_np.shape[0],
    }