VerificationCode/predict_model.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
# @Author  : bidikeji
# @Time    : 2019/11/25 0025 9:54 
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

from tensorflow.keras import models
import tensorflow.keras.backend as K
import tensorflow as tf
from PIL import Image
import numpy as np
import string

global graph
total_num = 0
neg_num = 0
graph = tf.get_default_graph()

digit_characters = string.digits
digit_base_model = models.load_model('gru_digit_base_model.h5')
arith_characters = '0123456789+*-%'
arith_base_model = models.load_model('gru_arith_base_model.h5')

def predict_digit(img):
    img_arr = np.array(img.resize((100, 50), Image.BILINEAR)) / 255.0
    X_test = np.array([img_arr])
    with graph.as_default():
        y_pred = digit_base_model.predict(X_test)
        out_pre = K.get_value(K.ctc_decode(y_pred, input_length=np.ones(y_pred.shape[0]) * y_pred.shape[1])[0][0])[:, :6]
    out = ''.join([digit_characters[x] for x in out_pre[0]])
    return out


def predict_arith(img):
    img_arr = np.array(img.resize((100, 50), Image.BILINEAR)) / 255.0
    X_test = np.array([img_arr])
    with graph.as_default():
        y_pred = arith_base_model.predict(X_test)
        out_pre = K.get_value(K.ctc_decode(y_pred, input_length=np.ones(y_pred.shape[0]) * y_pred.shape[1])[0][0])[:, :6]
    out = ''.join([arith_characters[x] for x in out_pre[0]])
    return out