fmard/ocr-api

1"""2OCR Service — PaddleOCR wrapper as HTTP API.3 4This service runs as a separate Python process and is called5by the Go orchestrator's OCR agent via HTTP.6 7Endpoints:8    POST /ocr  — Extract text from image9    GET  /health — Health check10"""11 12import io13import logging14import os15import time16 17import numpy as np18from flask import Flask, request, jsonify19from PIL import Image20from paddleocr import PaddleOCR21 22from preprocessor import preprocess23 24# Setup logging25logging.basicConfig(26    level=logging.INFO,27    format="%(asctime)s [%(levelname)s] %(message)s"28)29logger = logging.getLogger(__name__)30 31app = Flask(__name__)32 33# Multiple OCR engines for retry strategy34ocr_engines = {}35 36 37def get_ocr_engine(config_name="default"):38    """Lazy-load OCR engine with specific config."""39    global ocr_engines40    if config_name not in ocr_engines:41        logger.info(f"Loading PaddleOCR model config '{config_name}'...")42 43        configs = {44            "default": {45                "use_angle_cls": True,46                "lang": "en",47                "show_log": False,48                "use_gpu": False,49                "det_db_thresh": 0.2,       # Lower = detect more text regions50                "det_db_box_thresh": 0.3,    # Lower = keep more boxes51                "det_db_unclip_ratio": 1.8,  # Higher = bigger text boxes52                "rec_batch_num": 6,53            },54            "aggressive": {55                "use_angle_cls": True,56                "lang": "en",57                "show_log": False,58                "use_gpu": False,59                "det_db_thresh": 0.15,       # Very sensitive detection60                "det_db_box_thresh": 0.2,    # Keep almost all boxes61                "det_db_unclip_ratio": 2.0,  # Larger text regions62                "rec_batch_num": 6,63                "det_limit_side_len": 1920,  # Allow larger images64            },65        }66 67        cfg = configs.get(config_name, configs["default"])68        ocr_engines[config_name] = PaddleOCR(**cfg)69        logger.info(f"PaddleOCR model '{config_name}' loaded.")70 71    return ocr_engines[config_name]72 73 74@app.route("/health", methods=["GET"])75def health():76    """Health check endpoint."""77    return jsonify({"status": "ok", "engine": "paddleocr"})78 79 80@app.route("/ocr", methods=["POST"])81def ocr_extract():82    """Extract text from an uploaded image with multi-strategy OCR."""83    start_time = time.time()84 85    # Read image from request86    image_data = _read_image_from_request(request)87    if image_data is None:88        return jsonify({89            "success": False,90            "error": {"code": "no_image", "message": "No image provided"}91        }), 40092 93    try:94        # Preprocess image95        preprocess_meta = None96        image_np = None97 98        try:99            image_np, preprocess_meta = preprocess(image_data)100            logger.info(101                f"Preprocessing: blur={preprocess_meta['blur_score']:.3f}, "102                f"rotated={preprocess_meta['rotated']}, "103                f"cropped={preprocess_meta['cropped']}, "104                f"size={preprocess_meta['processed_size']}"105            )106        except Exception as preprocess_err:107            logger.warning(f"Preprocessing failed: {preprocess_err}")108 109        # Fallback: raw image decode if preprocessing failed110        if image_np is None:111            image = Image.open(io.BytesIO(image_data)).convert("RGB")112            image_np = np.array(image)113 114        # Strategy 1: Default OCR115        text_blocks = _run_ocr(image_np, "default")116 117        # Strategy 2: If too few results, try aggressive config118        if len(text_blocks) < 5:119            logger.info(120                f"Only {len(text_blocks)} blocks with default, "121                f"trying aggressive config..."122            )123            aggressive_blocks = _run_ocr(image_np, "aggressive")124            if len(aggressive_blocks) > len(text_blocks):125                text_blocks = aggressive_blocks126                logger.info(f"Aggressive config got {len(text_blocks)} blocks")127 128        # Strategy 3: If still few results, try without preprocessing (raw image)129        if len(text_blocks) < 5 and preprocess_meta is not None:130            logger.info("Trying OCR on raw image without preprocessing...")131            raw_image = Image.open(io.BytesIO(image_data)).convert("RGB")132            raw_np = np.array(raw_image)133            raw_blocks = _run_ocr(raw_np, "aggressive")134            if len(raw_blocks) > len(text_blocks):135                text_blocks = raw_blocks136                logger.info(f"Raw image got {len(text_blocks)} blocks")137 138        # Build response139        full_text_parts = [b["text"] for b in text_blocks]140        processing_time_ms = int((time.time() - start_time) * 1000)141 142        response = {143            "success": True,144            "data": {145                "text_blocks": text_blocks,146                "full_text": "\n".join(full_text_parts),147                "block_count": len(text_blocks),148                "processing_time_ms": processing_time_ms,149                "preprocessing": preprocess_meta,150            }151        }152 153        logger.info(154            f"OCR completed: {len(text_blocks)} blocks, "155            f"{processing_time_ms}ms"156        )157 158        return jsonify(response)159 160    except Exception as e:161        logger.error(f"OCR processing error: {e}")162        return jsonify({163            "success": False,164            "error": {"code": "processing_error", "message": str(e)}165        }), 500166 167 168def _run_ocr(image_np, config_name="default", max_retries=3):169    """Run PaddleOCR with retry + delay for memory recovery."""170    import gc171 172    for attempt in range(max_retries):173        try:174            engine = get_ocr_engine(config_name)175            results = engine.ocr(image_np, cls=True)176 177            text_blocks = []178            if results and results[0]:179                for line in results[0]:180                    bbox = line[0]181                    text = line[1][0]182                    confidence = float(line[1][1])183 184                    # Skip very low confidence results185                    if confidence < 0.3:186                        continue187 188                    text_blocks.append({189                        "text": text,190                        "confidence": confidence,191                        "bounding_box": {192                            "top_left": [int(bbox[0][0]), int(bbox[0][1])],193                            "top_right": [int(bbox[1][0]), int(bbox[1][1])],194                            "bottom_right": [int(bbox[2][0]), int(bbox[2][1])],195                            "bottom_left": [int(bbox[3][0]), int(bbox[3][1])],196                        }197                    })198 199            return text_blocks200 201        except Exception as e:202            logger.warning(f"OCR engine '{config_name}' attempt {attempt+1}/{max_retries} failed: {e}")203            # Force garbage collection and wait before retry204            gc.collect()205            if attempt < max_retries - 1:206                time.sleep(1)  # Wait 1s for memory to free up207 208    return []209 210 211def _read_image_from_request(req):212    """Extract image bytes from Flask request."""213    if "image" in req.files:214        return req.files["image"].read()215    if req.content_length and req.content_length > 0:216        return req.get_data()217    return None218 219 220if __name__ == "__main__":221    port = int(os.environ.get("OCR_SERVICE_PORT", 5001))222    logger.info(f"Starting OCR service on port {port}")223 224    # Pre-load default model on startup225    if os.environ.get("PRELOAD_MODEL", "true").lower() == "true":226        get_ocr_engine("default")227 228    app.run(host="0.0.0.0", port=port, debug=False)229