VerificationCode/chinese_equation_recognize/cer_interface_torch.py

import base64
import json
import logging
import os
import sys
import time
import traceback
from glob import glob
import cv2

# 只导入torch，protobuf会报错。需先导入TensorFlow再导入torch
import tensorflow
import torch
from flask import Flask, request
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
from chinese_equation_recognize.inference_equation_torch import recognize
from model_torch import crnn_ctc_equation_torch6
from utils import pil_resize, np2bytes, request_post, bytes2np, base64_decode, image_to_str, str_to_image

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

package_dir = os.path.abspath(os.path.dirname(__file__))
model_path = package_dir + "/models/equation6_model_acc-0.853.pth"
image_shape = (32, 192, 3)

# 接口配置
app = Flask(__name__)


@app.route('/cer', methods=['POST'])
def cer():
    start_time = time.time()
    logging.info("into cer_interface cer")
    try:
        # 接收网络数据
        if not request.form:
            logging.info("cer no data!")
            return json.dumps({"data": "", "success": 0})
        data = request.form.get("data")
        logging.info("cer_interface get data time" + str(time.time()-start_time))

        # 加载模型
        cer_model = globals().get("global_cer_model")
        if cer_model is None:
            print("=========== init cer model ===========")
            cer_model = CerModels().get_model()
            globals().update({"global_cer_model": cer_model})

        # 数据转换
        data = base64_decode(data)
        image_np = bytes2np(data)
        # 预测
        result = recognize(image_np, cer_model)
        if result is None:
            return json.dumps({"data": "", "success": 0})
        return json.dumps({"data": result, "success": 1})
    except:
        traceback.print_exc()
        return json.dumps({"data": "", "success": 0})
    finally:
        logging.info("cer interface finish time " + str(time.time()-start_time))


class CerModels:
    def __init__(self):
        device = torch.device("cpu")
        class_num = 35 + 1
        self.model = crnn_ctc_equation_torch6(class_num)
        self.model.load_state_dict(torch.load(model_path, map_location=torch.device(device)))
        self.model.eval()

    def get_model(self):
        return self.model


def test_cer_model(from_remote=True):
    paths = glob("D:/Project/captcha/data/test/FileInfo1021/1d419189-5116-11ed-851c-b4b5b67760ae_7.jpg")
    paths = glob(r'C:\Users\Administrator\Downloads\default.jfif')
    for file_path in paths:
        img_np = cv2.imread(file_path)
        h, w = img_np.shape[:2]
        file_bytes = np2bytes(img_np)
        file_base64 = base64.b64encode(file_bytes)

        if from_remote:
            file_json = {"data": file_base64}
            # _url = "http://192.168.2.102:17061/cer"
            _url = "http://127.0.0.1:17061/cer"
            result = json.loads(request_post(_url, file_json))
            if result.get("success"):
                result = int(result.get("data"))
                cv2.imshow("img_np", img_np)
                print("equation result", result)
                cv2.waitKey(0)
            else:
                print("failed!")


if __name__ == "__main__":
    # app.run(host='127.0.0.1', port=17061, debug=False)
    test_cer_model()