VerificationCode/chinese_recognize/chr_interface.py

import base64
import json
import logging
import os
import sys
import time
import traceback
from glob import glob
import cv2
import numpy as np
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
import tensorflow as tf
from flask import Flask, request
from chinese_recognize.model import cnn_net, cnn_net_tiny, cnn_net_small
from chinese_recognize.inference_char import recognize
from utils import pil_resize, np2bytes, request_post, bytes2np, base64_decode, str_to_image

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
tf.compat.v1.disable_eager_execution()
sess = tf.compat.v1.Session(graph=tf.Graph())

package_dir = os.path.abspath(os.path.dirname(__file__))
image_shape = (40, 40, 1)
model_path = package_dir + "/models/char_acc_0.89.h5"

# 接口配置
app = Flask(__name__)


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

        # 加载模型
        chr_model = globals().get("global_chr_model")
        if chr_model is None:
            print("=========== init chr model ===========")
            chr_model = ChrModels().get_model()
            globals().update({"global_chr_model": chr_model})

        # 数据转换
        str_list = json.loads(data)
        image_np_list = []
        for _str in str_list:
            image_np = str_to_image(_str)
            # b64 = _str.encode("utf-8")
            # image_np = bytes2np(base64_decode(b64))
            image_np_list.append(image_np)

        # 预测
        char_list = recognize(image_np_list, chr_model, sess)
        return json.dumps({"data": char_list, "success": 1})
    except:
        traceback.print_exc()
        return json.dumps({"data": "", "success": 0})
    finally:
        logging.info("chr interface finish time " + str(time.time()-start_time))


class ChrModels:
    def __init__(self):
        with sess.as_default():
            with sess.graph.as_default():
                self.model = cnn_net_small(input_shape=image_shape)
                self.model.load_weights(model_path)

    def get_model(self):
        return self.model


def test_chr_model(from_remote=True):
    paths = glob("D:/Project/captcha/data/test/char_9.jpg")
    str_list = []
    for file_path in paths:
        img_np = cv2.imread(file_path)
        cv2.imshow("img_np", img_np)
        cv2.waitKey(0)
        h, w = img_np.shape[:2]
        file_bytes = np2bytes(img_np)
        file_base64 = base64.b64encode(file_bytes)
        file_str = file_base64.decode("utf-8")
        str_list.append(file_str)
    if from_remote:
        file_json = {"data": json.dumps(str_list)}
        _url = "http://192.168.2.103:17057/chr"
        result = json.loads(request_post(_url, file_json))
        if result.get("success"):
            char_list = result.get("data")
            for i in range(len(paths)):
                print("image_path, char", paths[i], char_list[i])
        else:
            print("failed!")


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