from fastapi import FastAPI, File, UploadFile
from PIL import Image
from fastapi.middleware.cors import CORSMiddleware
from rapidocr_onnxruntime import RapidOCR
import io
import numpy as np
import pandas as pd

model = RapidOCR()
app = FastAPI()

origins = [
    "https://hycjack-fastapi-rapidocr.hf.space/",
    "http://localhost",
    "http://localhost:7860",
    "http://127.0.0.1:7860"
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"] 
)

@app.get("/")
def read_root():
    return {"Hello": "World!"}
    
@app.post("/ocr")    
async def ocr(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))
    np_array = np.array(image)
    ocr_result, elapse = model(np_array)
    dt_boxes, rec_res, scores = list(zip(*ocr_result))
    out_df = pd.DataFrame(
        [[box, rec, score] for box, rec, score in zip(dt_boxes, rec_res, scores)],
        columns=("box", "rec", "score"),
    )
    return out_df.to_dict(orient='records')