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()