ek-developers/my-ocr-api
0
1from fastapi import FastAPI, HTTPException, UploadFile, File, Form2from fastapi.responses import JSONResponse3from contextlib import asynccontextmanager4import httpx5import base646import os7import asyncio8import logging9 10logging.basicConfig(level=logging.INFO)11logger = logging.getLogger("ocr-api")12 13MODEL = "deepseek-ocr"14OLLAMA_URL = "http://localhost:11434"15 16# ── Tuned for free HuggingFace tier (2 vCPU, 16GB RAM, no GPU) ───17# OLLAMA_OPTIONS = {18# "temperature": 0, # deterministic — no sampling overhead19# "num_predict": 600, # enough for a simple document/image20# "num_thread": 2, # exactly match the 2 vCPUs21# "num_ctx": 512, # OCR doesn't need large context window22# "num_batch": 128, # smaller batch = less RAM pressure23# }24# OLLAMA_OPTIONS = {25# "temperature": 0, # deterministic26# "num_predict": 600, # enough for a simple document/image27# "num_thread": 2, # match 2 vCPUs28# "num_ctx": 512,29# "num_batch": 128,30# "repeat_penalty": 1.3, # prevents repetition loops31# "repeat_last_n": 64, # look back 64 tokens for repeats32# }33OLLAMA_OPTIONS = None34KEEP_ALIVE = "20m" # HF sleeps after 15 min anyway35 36 37# ─────────────────────────────────────────38# Warmup on startup39# ─────────────────────────────────────────40@asynccontextmanager41async def lifespan(app):42 asyncio.create_task(_warmup())43 yield44 45 46async def _warmup():47 """Pre-load model into RAM so first real request skips the 8s load time."""48 await asyncio.sleep(3)49 for attempt in range(5):50 try:51 async with httpx.AsyncClient(timeout=httpx.Timeout(60, connect=5, read=60)) as client:52 resp = await client.post(53 f"{OLLAMA_URL}/api/generate",54 json={55 "model": MODEL,56 "prompt": "hi",57 "stream": False,58 "keep_alive": KEEP_ALIVE,59 "options": OLLAMA_OPTIONS,60 }61 )62 if resp.status_code == 200:63 logger.info("Model warm and ready")64 return65 except Exception as e:66 logger.warning(f"Warmup attempt {attempt + 1}/5 failed: {e}")67 await asyncio.sleep(5)68 logger.warning("Warmup failed — model will load on first request")69 70 71app = FastAPI(lifespan=lifespan)72 73 74# ─────────────────────────────────────────75# Routes76# ─────────────────────────────────────────77@app.get("/")78def root():79 return {"status": "ok", "model": MODEL}80 81@app.head("/")82def head():83 return JSONResponse({})84 85@app.get("/health")86async def health():87 try:88 async with httpx.AsyncClient(timeout=3) as client:89 r = await client.get(f"{OLLAMA_URL}/api/tags")90 models = [m["name"] for m in r.json().get("models", [])]91 loaded = any(MODEL.split(":")[0] in m for m in models)92 return {"status": "ok", "model_loaded": loaded}93 except Exception as e:94 return {"status": "degraded", "error": str(e)}95 96 97@app.post("/ocr")98async def ocr(99 file: UploadFile = File(...),100 prompt_text: str = Form(101 default=(102 "Read all text visible in this image. "103 "Output only the raw text, word for word. "104 "No descriptions. No commentary. Just the text."105 )106 ),107 check_hallucination: bool = False,108 resize_image: bool = False,109 img_max_px: int = 1024,110):111 if not file.content_type or not file.content_type.startswith("image/"):112 raise HTTPException(400, "File must be an image")113 114 contents = await file.read()115 116 if resize_image:117 # if len(contents) > 5 * 1024 * 1024:118 # raise HTTPException(400, "Image too large. Maximum 5MB.")119 contents = _resize(contents,img_max_px)120 121 image_b64 = base64.b64encode(contents).decode()122 123 try:124 async with httpx.AsyncClient(125 timeout=httpx.Timeout(connect=10, read=180, write=30, pool=5)126 ) as client:127 resp = await client.post(128 f"{OLLAMA_URL}/api/generate",129 json={130 "model": MODEL,131 "prompt": prompt_text,132 "images": [image_b64],133 "stream": False,134 "keep_alive": KEEP_ALIVE,135 "options": OLLAMA_OPTIONS,136 }137 )138 resp.raise_for_status()139 text = resp.json().get("response", "").strip()140 141 if not text:142 raise HTTPException(500, "Model returned empty response")143 144 # If model described instead of transcribed, retry with stricter prompt145 if check_hallucination and _is_hallucinating(text):146 logger.warning("Model hallucinated — retrying with stricter prompt")147 resp2 = await client.post(148 f"{OLLAMA_URL}/api/generate",149 json={150 "model": MODEL,151 "prompt": "List every word and number printed on this document. Nothing else.",152 "images": [image_b64],153 "stream": False,154 "keep_alive": KEEP_ALIVE,155 "options": OLLAMA_OPTIONS,156 }157 )158 text = resp2.json().get("response", "").strip() or text159 160 logger.info(f"OCR complete — {len(text)} chars")161 return {"text": text}162 163 except httpx.TimeoutException:164 raise HTTPException(504, "OCR timed out. Try a clearer, smaller photo.")165 except httpx.HTTPStatusError as e:166 raise HTTPException(500, f"Model error: {e.response.text[:200]}")167 except HTTPException:168 raise169 except Exception as e:170 raise HTTPException(500, str(e))171 172 173# ─────────────────────────────────────────174# Helpers175# ─────────────────────────────────────────176def _resize(image_bytes: bytes, max_px: int = 1024) -> bytes:177 """Resize image to max 1024px on longest side, convert to JPEG."""178 try:179 from PIL import Image180 import io181 img = Image.open(io.BytesIO(image_bytes))182 if max(img.size) > max_px:183 img.thumbnail((max_px, max_px), Image.LANCZOS)184 buf = io.BytesIO()185 img.convert("RGB").save(buf, format="JPEG", quality=88, optimize=True)186 result = buf.getvalue()187 logger.info(f"Image resized: {len(image_bytes):,} → {len(result):,} bytes ({img.size})")188 return result189 except Exception as e:190 logger.warning(f"Resize failed ({e}) — using original")191 return image_bytes192 193 194def _is_hallucinating(text: str) -> bool:195 signals = [196 "the image", "the document", "appears to", "is presented",197 "is aligned", "without any formatting", "the text contains",198 "i can see", "this is a", "the card shows",199 ]200 return sum(1 for s in signals if s in text.lower()) >= 2