""" Kiri OCR - FastAPI OCR API """ import io import os import cv2 import tempfile import uvicorn import numpy as np from PIL import Image from fastapi import ( FastAPI, UploadFile, File, Form ) from fastapi.responses import ( JSONResponse, HTMLResponse ) from fastapi.middleware.cors import CORSMiddleware # ========================================================= # GLOBAL OCR INSTANCES # ========================================================= ocr_instances = {} # ========================================================= # LOAD OCR MODEL # ========================================================= def load_ocr(decode_method="accurate"): from kiri_ocr import OCR print( f"Loading OCR model with " f"decode_method={decode_method}" ) return OCR( model_path="mrrtmob/kiri-ocr", det_method="db", decode_method=decode_method, device="cpu", verbose=False ) # ========================================================= # GET OCR INSTANCE # ========================================================= def get_ocr(decode_method="accurate"): global ocr_instances if decode_method not in ocr_instances: ocr_instances[decode_method] = ( load_ocr(decode_method) ) return ocr_instances[decode_method] # ========================================================= # FASTAPI APP # ========================================================= app = FastAPI( title="Kiri OCR API", description="Image OCR API using Kiri OCR", version="1.0" ) # ========================================================= # CORS # ========================================================= app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ========================================================= # HOME PAGE # ========================================================= @app.get("/") async def home(): return HTMLResponse(""" Kiri OCR API

Kiri OCR API

Upload image and extract text 



        


        

    

    

    """)

# =========================================================
# OCR API
# =========================================================

@app.post("/img2ocr")
async def img2ocr(
    file: UploadFile = File(...),
    decode_method: str = Form("accurate")
):

    try:

        contents = await file.read()

        image = Image.open(
            io.BytesIO(contents)
        ).convert("RGB")

        ocr_engine = get_ocr(decode_method)

        image_np = np.array(image)

        image_bgr = cv2.cvtColor(
            image_np,
            cv2.COLOR_RGB2BGR
        )

        with tempfile.NamedTemporaryFile(
            suffix=".png",
            delete=False
        ) as f:

            temp_path = f.name

        cv2.imwrite(
            temp_path,
            image_bgr
        )

        extracted_text = ""

        # OCR STREAMING
        for chunk in (
            ocr_engine.extract_text_stream_chars(
                temp_path,
                mode="lines"
            )
        ):

            token = chunk.get(
                "token",
                ""
            )

            if token:
                extracted_text += token

        # CLEANUP
        if os.path.exists(temp_path):
            os.unlink(temp_path)

        return JSONResponse({

            "success": True,

            "decode_method":
                decode_method,

            "text":
                extracted_text

        })

    except Exception as e:

        return JSONResponse({

            "success": False,

            "error": str(e)

        })

# =========================================================
# HEALTH CHECK
# =========================================================

@app.get("/health")
async def health():

    return {
        "status": "running"
    }

# =========================================================
# MAIN
# =========================================================

if __name__ == "__main__":

    uvicorn.run(
        app,
        host="0.0.0.0",
        port=7860
    )