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

from core.direction import AngleDetector
from sx_utils.sximage import *
from sx_utils.sxtime import sxtimeit
from sx_utils.sxweb import web_try
from core.ocr import CetOcr
import os

# 导入一些包


app = FastAPI()

origins = ["*"]

# CORS 跨源资源共享
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
# templates = Jinja2Templates(directory='templates')

use_gpu = False
if os.getenv('USE_CUDA') == 'gpu':
    use_gpu = True

print(f'use gpu: {use_gpu}')

# 初始化ocr模型和后处理模型
# 分类
# ocr = PaddleOCR(use_angle_cls=True,
#                 # 方向
#                 rec_model_dir="./idcard_rec_infer/",
#                 det_model_dir="./idcard_det_infer/",
#                 cls_model_dir="idcard_cls_infer",
#                 # 识别
#                 rec_algorithm='CRNN',
#                 ocr_version='PP-OCRv2',
#                 # 中文字典
#                 rec_char_dict_path="./ppocr_keys_v1.txt", lang="ch",
#                 use_gpu=use_gpu,
#                 # 预训练-->效果不明显
#                 # 网络不够大、不够深
#                 # 数据集普遍较小，batch size普遍较小
#                 warmup=True)
# ocr = PaddleOCR(use_angle_cls=True,
#                 use_gpu=use_gpu)


ocr = PaddleOCR(use_angle_cls=True,
                use_gpu=use_gpu,
                det_db_unclip_ratio=2.5,
                det_db_thresh=0.1,
                det_db_box_thresh=0.3,
                warmup=True)
#
# ocr = PaddleOCR(use_angle_cls=True,
#                 rec_model_dir='./ch_ppocr_server_v2.0_rec_infer',
#                 det_model_dir='./ch_ppocr_server_v2.0_det_infer',
#                 cls_model_dir='./idcard_cls_infer',
#                 ocr_version='PP-OCRv2',
#                 rec_algorithm='CRNN',
#                 use_gpu=use_gpu,
#                 det_db_unclip_ratio=2.5,
#                 det_db_thresh=0.1,
#                 det_db_box_thresh=0.3,
#                 warmup=True)


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

# 初始化 身份证ocr识别 对象
m = CetOcr(ocr, ad)


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


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


# /ocr_system/bankcard 银行卡
# /ocr_system/regbook 户口本
# /ocr_system/schoolcert 学信网

# Post 接口
# 计算耗时
# 异常处理
@app.post("/ocr_system/cet")
@sxtimeit
@web_try()
# 传入=> base64码 -> np
# 返回=> 检测到到结果 -> (conf, angle, parser, image_type)
def cet(request: Request, cer: CetInfo):
    image = base64_to_np(cer.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)