Asib27's picture
add model code
c2bf863
raw
history blame
1.7 kB
import gradio as gr
import numpy as np
import cv2
import tensorflow as tf
import keras
from keras import layers, models
model = model = tf.keras.models.load_model('model/ocr_model.h5')
def preprocessImage(img, shape):
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img = cv2.resize(img, (shape))
img = (img/255).astype(np.float32)
img = img.T
img = np.expand_dims(img, axis=-1)
return img
label2char ={0: ' ',
1: "'",
2: '-',
3: 'A',
4: 'B',
5: 'C',
6: 'D',
7: 'E',
8: 'F',
9: 'G',
10: 'H',
11: 'I',
12: 'J',
13: 'K',
14: 'L',
15: 'M',
16: 'N',
17: 'O',
18: 'P',
19: 'Q',
20: 'R',
21: 'S',
22: 'T',
23: 'U',
24: 'V',
25: 'W',
26: 'X',
27: 'Y',
28: 'Z',
29: '`'}
def getStringFromEncode(lst :list):
return ''.join([label2char[i] if i in label2char else '' for i in lst])
def decode_batch_predictions(pred):
pred = pred[:, :-2] # first two layers of ctc garbage
input_len = np.ones(pred.shape[0])*pred.shape[1]
results = keras.backend.ctc_decode(pred,
input_length=input_len,
greedy=True)[0][0]
output_text = []
for res in results.numpy():
outstr = getStringFromEncode(res)
output_text.append(outstr)
# return final text results
return output_text
def predict(img):
img = preprocessImage(img, (256,64))
img = np.expand_dims(img, axis=0) # 1 image in batch
preds = model.predict(img)
pred_texts = decode_batch_predictions(preds)
return pred_texts[0]
def greet(name):
return "Hello " + name + "!!"
iface = gr.Interface(fn=greet, inputs="text", outputs="text")
iface.launch()