hr-ocr-idcard/server.py

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
from paddleocr import PaddleOCR
from core.idcrad import IdCardStraight
from core.direction import *
from base64 import b64decode
import numpy as np
from sx_utils.sximage import *
from sx_utils.sxtime import sxtimeit
from sx_utils.sxweb import web_try
import cpca



app = FastAPI()

origins = [
    "*"
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
# templates = Jinja2Templates(directory='templates')


# 初始化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=False)



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

# @app.get("/")
# def home(request: Request):
#     ''' Returns html jinja2 template render for home page form
#     '''
#
#     return templates.TemplateResponse('home.html', {
#             "request": request,
#         })


class IdCardInfo(BaseModel):
    image: str
    image_type: str


@app.post("/ocr_system/idcard")
@sxtimeit
@web_try()
def idcard(request: Request, id_card: IdCardInfo):
    image = base64_to_np(id_card.image)
    if int(id_card.image_type) == 0:
        return _front(image, id_card.image_type)
    if int(id_card.image_type) == 1:
       return _back(image, id_card.image_type)

def _back(image, image_type):
    raise Exception('not implemented yet')

def _front(image, image_type: str):
    angle = detect_angle(image)
    print(angle)
    # 获取模型检测结果
    result = ocr.ocr(image, cls=True)
    print("------------------")
    print(result)
    if not result:
        return None
    scores = [line[1][1] for line in result]

    score = sum(scores) / len(scores)
    print("------------------")
    print(scores)
    print("------------------")

    sc = []
    for i in range(0, 6):
        sc.append(scores[i])
    sc.append(scores[i])
    sc.append(scores[i])
    sc.append(scores[i])

    # 将检测到的文字放到一个列表中
    txts = [line[1][0] for line in result]
    print("......................................")
    print(txts)
    print("......................................")
    # 将结果送入到后处理模型中
    postprocessing = IdCardStraight(txts)
    id_result = postprocessing.run()
    content = id_result['Data']['FrontResult']

    location_str = []
    location_str.append(content["Address"])
    print(location_str)
    df = cpca.transform(location_str)
    print(df)

    province = df.iloc[0, 0]
    city = df.iloc[0, 1]
    region = df.iloc[0, 2]
    detail = df.iloc[0, 3]
    print(f'pronvince: {province}, city: {city}, region: {region}, detail: {detail}')
    return {
        "confidence": score,
        "card_type": image_type,
        "orientation": angle // 90,
        "name": {'text': content['Name'], 'confidence': sc[0]},
        "id": {'text': content['IDNumber'], 'confidence': sc[1]},
        "ethnicity": {'text': content['Nationality'], 'confidence': sc[2]},
        "gender": {'text': content['Gender'], 'confidence': sc[3]},
        "birthday": {'text': content['year'], 'confidence': sc[4]},
        "address_province": {'text': province, 'confidence': sc[5]},
        "address_city": {'text': city, 'confidence': sc[6]},
        "address_region": {'text': region, 'confidence': sc[7]},
        "address_detail": {'text': detail, 'confidence': sc[8]},
        "expire_date": {'text': '', 'confidence': 0}
    }





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)