import os os.environ["OMP_NUM_THREADS"] = "1" os.environ["MKL_NUM_THREADS"] = "1" import torch # PREVENT MEMORY ISSUES IN DOCKER: # Set torch thread limits before ANY other library initialization. try: torch.set_num_threads(1) torch.set_num_interop_threads(1) except RuntimeError: # Already initialized, that's fine pass import os, base64, traceback, requests, uuid, io from fastapi import FastAPI from pydantic import BaseModel from typing import Optional from PIL import Image from src.core.aadhaar import process_image from src.core.pan import process_pan from src.utils.ocr_engines import get_easyocr_reader app = FastAPI() # Warmup OCR reader at startup @app.on_event("startup") async def startup_event(): print("Warming up OCR engines...") get_easyocr_reader() print("OCR engines warmed up and ready.") UPLOAD_DIR = "ocr_uploads" os.makedirs(UPLOAD_DIR, exist_ok=True) class OCRRequest(BaseModel): url: Optional[str] = None fileB64: Optional[str] = None id: int user_id: int file_type: str @app.post("/ocr/aadhaar_ocr") async def ocr_endpoint(payload: OCRRequest): file_path = None try: extension = ".bin" # ======================== # Base64 Input # ======================== if payload.fileB64: b64_str = payload.fileB64.strip() # Detect MIME from data URI if b64_str.startswith("data:"): header, b64_str = b64_str.split(",", 1) if "jpeg" in header or "jpg" in header: extension = ".jpg" elif "png" in header: extension = ".png" elif "pdf" in header: extension = ".pdf" elif "webp" in header: extension = ".webp" else: # Auto detect from Base64 signature if b64_str.startswith("JVBERi0"): extension = ".pdf" elif b64_str.startswith("/9j/"): extension = ".jpg" elif b64_str.startswith("iVBORw0"): extension = ".png" elif b64_str.startswith("UklGR"): extension = ".webp" file_bytes = base64.b64decode(b64_str) unique_name = str(uuid.uuid4()) # Convert WEBP to JPG for better OCR support if extension == ".webp": image = Image.open(io.BytesIO(file_bytes)) file_path = os.path.join(UPLOAD_DIR, f"{unique_name}.jpg") image.convert("RGB").save(file_path, "JPEG") else: file_path = os.path.join(UPLOAD_DIR, f"{unique_name}{extension}") with open(file_path, "wb") as f: f.write(file_bytes) # ======================== # URL Input # ======================== elif payload.url: r = requests.get(payload.url, timeout=15) if r.status_code != 200: return { "status": "failed", "result": [], "message": "File download failed" } from urllib.parse import urlparse url_path = urlparse(payload.url).path extension = os.path.splitext(url_path)[1].lower() or ".bin" unique_name = str(uuid.uuid4()) file_bytes = r.content # Convert WEBP from URL too if extension == ".webp": image = Image.open(io.BytesIO(file_bytes)) file_path = os.path.join(UPLOAD_DIR, f"{unique_name}.jpg") image.convert("RGB").save(file_path, "JPEG") else: file_path = os.path.join(UPLOAD_DIR, f"{unique_name}{extension}") with open(file_path, "wb") as f: f.write(file_bytes) else: return { "status": "failed", "result": [], "message": "No fileB64 or url provided" } # ======================== # OCR Processing # ======================== response = await process_image( file_path, payload.id, payload.user_id, payload.file_type ) return response except Exception: error_msg = traceback.format_exc() print(f"Error in ocr_endpoint: {error_msg}") return { "status": "failed", "result": [], "message": error_msg } finally: # file cleaup if file_path and os.path.exists(file_path): os.remove(file_path) @app.post("/ocr/pan_ocr") async def pan_ocr_endpoint(payload: OCRRequest): file_path = None try: extension = ".bin" # ======================== # Base64 Input # ======================== if payload.fileB64: b64_str = payload.fileB64.strip() # Detect MIME from data URI if b64_str.startswith("data:"): header, b64_str = b64_str.split(",", 1) if "jpeg" in header or "jpg" in header: extension = ".jpg" elif "png" in header: extension = ".png" elif "pdf" in header: extension = ".pdf" elif "webp" in header: extension = ".webp" else: # Auto detect from Base64 signature if b64_str.startswith("JVBERi0"): extension = ".pdf" elif b64_str.startswith("/9j/"): extension = ".jpg" elif b64_str.startswith("iVBORw0"): extension = ".png" elif b64_str.startswith("UklGR"): extension = ".webp" file_bytes = base64.b64decode(b64_str) unique_name = str(uuid.uuid4()) # Convert WEBP to JPG for better OCR support if extension == ".webp": image = Image.open(io.BytesIO(file_bytes)) file_path = os.path.join(UPLOAD_DIR, f"{unique_name}.jpg") image.convert("RGB").save(file_path, "JPEG") else: file_path = os.path.join(UPLOAD_DIR, f"{unique_name}{extension}") with open(file_path, "wb") as f: f.write(file_bytes) # ======================== # URL Input # ======================== elif payload.url: r = requests.get(payload.url, timeout=15) if r.status_code != 200: return { "status": "failed", "result": [], "message": "File download failed" } from urllib.parse import urlparse url_path = urlparse(payload.url).path extension = os.path.splitext(url_path)[1].lower() or ".bin" unique_name = str(uuid.uuid4()) file_bytes = r.content # Convert WEBP from URL too if extension == ".webp": image = Image.open(io.BytesIO(file_bytes)) file_path = os.path.join(UPLOAD_DIR, f"{unique_name}.jpg") image.convert("RGB").save(file_path, "JPEG") else: file_path = os.path.join(UPLOAD_DIR, f"{unique_name}{extension}") with open(file_path, "wb") as f: f.write(file_bytes) else: return { "status": "failed", "result": [], "message": "No fileB64 or url provided" } # ======================== # OCR Processing # ======================== response = await process_pan( file_path, payload.id, payload.user_id, payload.file_type ) return response except Exception: error_msg = traceback.format_exc() print(f"Error in pan_ocr_endpoint: {error_msg}") return { "status": "failed", "result": [], "message": error_msg } finally: # file cleaup if file_path and os.path.exists(file_path): os.remove(file_path) if __name__ == "__main__": import uvicorn uvicorn.run( "app_ocr:app",host="127.0.0.1", port=5001, reload=True, log_level="debug" )