ocr-orderid3 / wapp.py
ethanrom's picture
Duplicate from ethanrom/ocr-orderid
7d30f01
import cv2
import pytesseract
import streamlit as st
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import load_img, img_to_array
def main():
st.set_page_config(page_title='Order ID Finder', layout='wide')
st.title('OCR + Font type demo')
st.write('uses OCR to extract text from an uploaded image of a name. The app then attempts to match the extracted text with the corresponding order ID from a uploaded text file. Also uses pre-trained TensorFlow CNN to classify the font type of the extracted text as either Pacifico or OpenSans-Light. If the app successfully matches the text with an order ID, it outputs the order ID, predicted font type, and detected text. If no match is found, it outputs the predicted font type and detected text.')
with st.sidebar:
st.write('## Input')
uploaded_file = st.file_uploader('Upload the image file (PNG or JPG)', type=['png', 'jpg'], help='help')
input_file = st.file_uploader('Upload the input file (TXT)', type=['txt'], help='help')
if st.button('Find Order ID') and uploaded_file and input_file:
st.write('## Output')
img = cv2.imdecode(np.fromstring(uploaded_file.read(), np.uint8), 1)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
text = pytesseract.image_to_string(gray)
with input_file as file:
file_contents = file.read().decode()
lines = file_contents.split('\n')
found = False
for line in lines:
order_id, name, font = line.strip().split(',')
if name.strip() in text:
image = load_img(uploaded_file, target_size=(64, 64))
image = img_to_array(image)
image = np.expand_dims(image, axis=0)
image = image / 255.0
prediction = model.predict(image)
font_type = 'Pacifico' if prediction[0, 0] > prediction[0, 1] else 'OpenSans-Light'
st.success(f'Detected Text: {text.strip()}\n, Order ID: {order_id}, Predicted Font Type: {font_type}')
found = True
break
if not found:
image = load_img(uploaded_file, target_size=(64, 64))
image = img_to_array(image)
image = np.expand_dims(image, axis=0)
image = image / 255.0
prediction = model.predict(image)
font_type = 'Pacifico' if prediction[0, 0] > prediction[0, 1] else 'OpenSans-Light'
st.warning(f'Detected Text: {text.strip()}\n, Could not find the Order ID, Predicted Font Type: {font_type}')
if __name__ == '__main__':
model = tf.keras.models.load_model('model.h5')
main()