from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from paddleocr import PaddleOCR

from blfe_core.direction import AngleDetector
from utils.image import *
from utils.time import timeit
from utils.web import web_try
from blfe_core.ocr import BusinessLicenseOcr
import os

app = FastAPI()

origins = ["*"]

# CORS 跨源资源共享
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
use_gpu = os.getenv('USE_CUDA') == 'gpu'
print(f'use gpu: {use_gpu}')

ocr = PaddleOCR(
                det_model_dir="./models/det",
                rec_model_dir="./models/rec",
                cls_model_dir="./models/cls",
                use_gpu=use_gpu,
                det_db_unclip_ratio=1.8,
                det_db_thresh=0.1,
                det_db_box_thresh=0.3,)

# 初始化 角度检测器 对象
ad = AngleDetector(ocr)

# 初始化 ocr识别 对象
m = BusinessLicenseOcr(ocr, ad)


# Get 健康检查
@app.get("/ping")
def ping():
    return "pong!"


# 解析传入的 json对象
class BusinessLicenseInfo(BaseModel):
    image: str


# Post 接口
@app.post("/ocr_system/business_license")
@timeit
@web_try()
def blfe(request: Request, blfe: BusinessLicenseInfo):
    image = base64_to_np(blfe.image)
    return m.predict(image)


if __name__ == '__main__':
    import uvicorn
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--host', default='0.0.0.0')
    parser.add_argument('--port', default=8080)
    opt = parser.parse_args()

    app_str = 'server:app'  # make the app string equal to whatever the name of this file is
    uvicorn.run(app_str, host=opt.host, port=int(opt.port), reload=True)