Update app.py

app.py

@@ -1,134 +1,191 @@
 import os
-import asyncio
+import uuid
 import time
-import io
 import gc
+import asyncio
 from concurrent.futures import ProcessPoolExecutor
-from typing import List
-import numpy as np
-from PIL import Image
-from fastapi import FastAPI, UploadFile, File, HTTPException
+from typing import List, Dict, Any
+from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks
 from fastapi.responses import JSONResponse
-from paddleocr import PaddleOCR
+import threading
 
-# --- CPU-ONLY CONFIGURATION ---
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
-os.environ["PADDLE_HOME"] = "/app/paddle_home"
-os.makedirs("/app/paddle_home", exist_ok=True)
+# ───────────────────────────────────────────────────────────────────
+# CONFIGURATION (2 vCPU + 16GB RAM OPTIMIZED)
+# ───────────────────────────────────────────────────────────────────
+MAX_WORKERS = 3  # 2 cores + 1 overlap
+MAX_FILES = 15
+CACHE_TTL = 3600
 
-# --- WORKER INITIALIZATION ---
-def init_worker():
-    """Initialize OCR engine once per worker process"""
-    global worker_ocr
-    # Let PaddleOCR auto-select Marathi models - cleanest approach
-    worker_ocr = PaddleOCR(
-        lang="mr",
-        use_gpu=False,
-        enable_mkldnn=True,  # 2-3x speedup on modern CPUs
-        cpu_threads=2,       # Threads per worker (tune 1-3)
-    )
+os.environ.update({
+    "PADDLE_HOME": "/app/paddle_home", 
+    "XDG_CACHE_HOME": "/app/xdg_cache"
+})
 
-# --- RESOURCE MANAGEMENT ---
-MAX_WORKERS = min(4, max(1, os.cpu_count() - 1))
-CPU_EXECUTOR = ProcessPoolExecutor(max_workers=MAX_WORKERS, initializer=init_worker)
+for path in ["/app/paddle_home", "/app/xdg_cache", "/app/uploads"]:
+    os.makedirs(path, exist_ok=True)
 
-# --- FASTAPI APP ---
-app = FastAPI(title="CPU-Optimized OCR API", version="1.0.0")
+# ───────────────────────────────────────────────────────────────────
+# GLOBAL STATE
+# ───────────────────────────────────────────────────────────────────
+executor = ProcessPoolExecutor(max_workers=MAX_WORKERS)
+semaphore = asyncio.Semaphore(MAX_WORKERS)
+active_count = 0
+state_lock = threading.Lock()
 
-@app.middleware("http")
-async def add_metrics(request, call_next):
-    """Add performance metrics to response headers"""
-    start = time.perf_counter()
-    response = await call_next(request)
-    response.headers["X-Process-Time"] = f"{time.perf_counter() - start:.3f}s"
-    return response
+app = FastAPI(title="Ultra-Fast Image OCR API", version="3.0")
 
-# --- CORE PROCESSING ---
-def process_single_image(image_data: bytes) -> dict:
-    """Process one image in isolated worker process"""
+# ───────────────────────────────────────────────────────────────────
+# WORKER FUNCTION
+# ───────────────────────────────────────────────────────────────────
+def process_image(file_tuple: tuple) -> Dict[str, Any]:
+    """Process single image in isolated worker: (file_id, filepath)"""
+    # Lazy-load model once per worker process
+    if not hasattr(process_image, 'engine'):
+        from paddleocr import PaddleOCR
+        process_image.engine = PaddleOCR(
+            lang="mr",
+            text_recognition_model_name="devanagari_PP-OCRv5_mobile_rec",
+            use_doc_orientation_classify=False,
+            use_doc_unwarping=False,
+            use_textline_orientation=False
+        )
+    
+    file_id, filepath = file_tuple
+    
     try:
-        # Zero-copy image loading from memory
-        img = Image.open(io.BytesIO(image_data)).convert("RGB")
-        np_img = np.array(img)
-        
-        # Run OCR
-        ocr_result = worker_ocr.predict(input=np_img)
-        
-        # Parse results
-        results = [
-            {"text": text, "confidence": float(score)}
-            for block in ocr_result
-            for text, score in zip(block.get("rec_texts", []), block.get("rec_scores", []))
-        ]
+        start = time.perf_counter()
+        result = process_image.engine.predict(filepath)
         
-        return {"success": True, "results": results}
+        return {
+            "file_id": file_id,
+            "status": "success",
+            "processing_time_seconds": round(time.perf_counter() - start, 3),
+            "results": [
+                {"text": txt, "confidence": float(conf)}
+                for block in result
+                for txt, conf in zip(block["rec_texts"], block["rec_scores"])
+            ]
+        }
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {
+            "file_id": file_id,
+            "status": "error",
+            "processing_time_seconds": 0,
+            "error": str(e),
+            "results": []
+        }
     finally:
-        # Lightweight cleanup
-        gc.collect()
+        os.remove(filepath)
+
+# ───────────────────────────────────────────────────────────────────
+# BACKGROUND MAINTENANCE
+# ───────────────────────────────────────────────────────────────────
+def cleanup_cache():
+    """Remove old uploaded files"""
+    cutoff = time.time() - CACHE_TTL
+    for f in os.listdir("/app/uploads"):
+        try:
+            path = os.path.join("/app/uploads", f)
+            if os.path.getmtime(path) < cutoff:
+                os.remove(path)
+        except:
+            pass
 
+# ───────────────────────────────────────────────────────────────────
+# API ENDPOINTS
+# ───────────────────────────────────────────────────────────────────
 @app.post("/ocr")
-async def ocr_endpoint(files: List[UploadFile] = File(...)):
-    """Process up to 10 images in parallel"""
-    if len(files) > 10:
-        raise HTTPException(status_code=400, detail="Maximum 10 files allowed.")
+async def ocr_endpoint(
+    background_tasks: BackgroundTasks,
+    files: List[UploadFile] = File(...)
+):
+    """Process up to 15 images with parallelization and wait tracking"""
+    if len(files) > MAX_FILES:
+        raise HTTPException(400, f"Maximum {MAX_FILES} images allowed")
+    
+    # Validate image types only
+    for f in files:
+        if not f.filename.lower().endswith(('.jpg', '.jpeg', '.png')):
+            raise HTTPException(400, f"Unsupported: {f.filename}. Use JPG/PNG only.")
+    
+    background_tasks.add_task(cleanup_cache)
+    request_start = time.perf_counter()
+    
+    # Save files
+    file_tuples = []
+    for idx, file in enumerate(files, 1):
+        ext = file.filename.split('.')[-1].lower()
+        path = f"/app/uploads/{uuid.uuid4()}.{ext}"
+        with open(path, "wb") as f:
+            f.write(await file.read())
+        file_tuples.append((f"file_{idx}", path))
     
-    # --- VALIDATION & LOAD ---
-    file_data = []
-    for file in files:
-        # Check format
-        ext = file.filename.split(".")[-1].lower()
-        if ext not in ("jpg", "jpeg", "png"):
-            raise HTTPException(status_code=400, detail=f"Unsupported: {file.filename}")
+    # Check capacity before processing
+    with state_lock:
+        had_to_wait = active_count >= MAX_WORKERS
+    
+    # Process with concurrency control
+    async def track_and_process(file_info):
+        global active_count
+        file_id = file_info[0]
         
-        # Read file (max 10MB per file)
-        content = await file.read()
-        if len(content) > 10 * 1024 * 1024:
-            raise HTTPException(status_code=400, detail=f"File too large: {file.filename}")
+        with state_lock:
+            active_count += 1
         
-        file_data.append(content)
-    
-    # --- PARALLEL PROCESSING ---
-    loop = asyncio.get_event_loop()
-    tasks = [
-        loop.run_in_executor(CPU_EXECUTOR, process_single_image, data)
-        for data in file_data
-    ]
+        try:
+            async with semaphore:
+                loop = asyncio.get_event_loop()
+                return await loop.run_in_executor(executor, process_image, file_info)
+        finally:
+            with state_lock:
+                active_count -= 1
     
-    results = await asyncio.gather(*tasks)
+    results = await asyncio.gather(*[track_and_process(ft) for ft in file_tuples])
     
-    # --- BUILD RESPONSE ---
-    structured_output = {"files": []}
-    for idx, result in enumerate(results):
-        file_info = {
-            "file_id": f"file_{idx + 1}",
-            "filename": files[idx].filename,
+    response = {
+        "files": results,
+        "performance_summary": {
+            "total_request_time_seconds": round(time.perf_counter() - request_start, 3),
+            "successful_files": sum(1 for r in results if r["status"] == "success"),
+            "failed_files": sum(1 for r in results if r["status"] == "error"),
+            "had_to_wait_for_worker": had_to_wait,
+            "request_status": "processed_immediately" if not had_to_wait else "processed_after_waiting"
         }
-        if result["success"]:
-            file_info["pages"] = [{"page_index": 0, "results": result["results"]}]
-        else:
-            file_info["error"] = result["error"]
-        structured_output["files"].append(file_info)
+    }
     
-    # Final cleanup
-    del file_data
     gc.collect()
-    
-    return JSONResponse(structured_output)
+    return response
 
-@app.get("/health")
-async def health_check():
-    """System status endpoint"""
+@app.get("/status")
+async def get_status():
+    """Live capacity monitor - check this before sending requests"""
+    with state_lock:
+        current = active_count
+    
+    available = MAX_WORKERS - current
+    
     return {
-        "status": "healthy",
-        "cpu_workers": MAX_WORKERS,
-        "cpu_threads_per_worker": 2,
-        "supported_formats": ["jpg", "jpeg", "png"],
-        "max_file_size": "10MB",
+        "system_capacity": {
+            "currently_processing": current,
+            "max_capacity": MAX_WORKERS,
+            "available_slots": available
+        },
+        "can_accept_immediately": available > 0,
+        "recommended_action": "send_request" if available > 0 else "wait_or_retry_later"
     }
 
+@app.delete("/cache")
+async def clear_cache():
+    """Manual cache cleanup"""
+    cleanup_cache()
+    gc.collect()
+    return {"status": "cache_cleared"}
+
+@app.get("/health")
+async def health():
+    """Health check"""
+    return {"status": "healthy", "active_processes": active_count}
+
 @app.on_event("shutdown")
-async def shutdown_workers():
-    """Graceful shutdown"""
-    CPU_EXECUTOR.shutdown(wait=True)
+def shutdown():
+    executor.shutdown(wait=True)