VerificationCode/chinese_detect/chd_interface.py

import base64
import json
import logging
import os
import sys
import time
import traceback
from glob import glob
import numpy as np
import cv2
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_detect.inference_yolo_char import get_tiny_inference_model, detect
from utils import pil_resize, np2bytes, request_post, bytes2np, get_anchors, get_classes, get_colors

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__))
model_path = package_dir + "/models/char_yolo_loss_39.90.h5"
anchors = get_anchors(package_dir + "/yolo_data/my_anchors.txt")
classes = get_classes(package_dir + "/yolo_data/my_classes.txt")
colors = get_colors(len(classes))

# 接口配置
app = Flask(__name__)


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

        # 加载模型
        chd_model = globals().get("global_chd_model")
        if chd_model is None:
            print("=========== init chd model ===========")
            chd_model = ChdModels().get_model()
            globals().update({"global_chd_model": chd_model})

        # 数据转换
        data = base64.b64decode(data)
        image_np = bytes2np(data)
        # 预测
        _, out_boxes, out_classes = detect(image_np, chd_model, sess, is_tips=is_tips)
        return json.dumps({"data": out_boxes, "success": 1})
    except:
        traceback.print_exc()
        return json.dumps({"data": "", "success": 0})
    finally:
        logging.info("chd interface finish time " + str(time.time()-start_time))


class ChdModels:
    def __init__(self):
        # detect
        with sess.as_default():
            with sess.graph.as_default():
                self.model = get_tiny_inference_model(anchors, len(classes), weights_path=model_path)

    def get_model(self):
        return self.model


def test_chd_model(from_remote=True):
    paths = glob("D:/Project/captcha/data/test/phrase_5.jpg")
    for file_path in paths:
        img_np = cv2.imread(file_path)
        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:17000/ocr"
            _url = "http://127.0.0.1:17000/chd"
            result = json.loads(request_post(_url, file_json))
            if result.get("success"):
                out_boxes = result.get("data")
                print("out_boxes", out_boxes)
                for box in out_boxes:
                    cv2.rectangle(img_np, (box[0], box[1]), (box[2], box[3]), (0, 0, 255))
                cv2.imshow("img_np", img_np)
                cv2.waitKey(0)
            else:
                print("failed!")


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