hr-ocr-bankcard/server.py

import os
from typing import Optional

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

from core.ocr import BankOcr
from core.direction import AngleDetector
from s_utils.image import *
from s_utils.time import timeit
from s_utils.web import web_try

app = FastAPI()

origins = ["*"]

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/models_0/det",
                rec_model_dir="./models/models_0/rec",
                cls_model_dir="./models/models_0/cls",
                rec_char_dict_path="./ppocr_keys_bank.txt",
                det_db_unclip_ratio=2.5,
                ocr_version='PP-OCRv2',
                det_db_thresh=0.3,
                det_db_box_thresh=0.6,
                use_angle_cls=True,
                use_gpu=use_gpu,
                use_space_char=False,
                warmup=True
                )


origin_ocr = PaddleOCR(
                       det_model_dir="./models/models_1/det",
                       rec_model_dir="./models/models_0/rec",
                       cls_model_dir="./models/models_1/cls",
                       rec_char_dict_path="./ppocr_keys_bank.txt",
                       det_db_unclip_ratio=2.5,
                       ocr_version='PP-OCRv2',
                       det_db_thresh=0.3,
                       det_db_box_thresh=0.6,
                       use_angle_cls=True,
                       use_gpu=use_gpu,
                       use_space_char=False,
                       warmup=True
                       )


ad = AngleDetector(origin_ocr)
m = BankOcr(ocr, ad)


@app.get("/ping")
def ping():
    return "pong!"


class IdCardInfo(BaseModel):
    image: str
    fn: Optional[str]

@app.post("/ocr_system/bankcard")
@timeit
@web_try()
def detect(request: Request, id_card: IdCardInfo):
    image = base64_to_np(id_card.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, workers=1)