abspython's picture
OCR to UI
b4dae5d
Raw
History Blame
2.25 kB
import streamlit as st
import numpy as np
import joblib
import easyocr
from transformers import pipeline
from keras.utils import load_img
from keras.utils import img_to_array
from PIL import Image
import io
from tempfile import NamedTemporaryFile
st.set_option('deprecation.showfileUploaderEncoding', False)
def load_image(image_file):
img = Image.open(image_file)
return img
def get_img_prediction(imgpath):
img = load_img(imgpath,target_size=(128,128,3))
img = img_to_array(img)
img = img/255
X_pred_image = np.array(img)
X_pred_imaged = X_pred_image.reshape(1,128*128*3)
y_pred_pro = loaded_lgbm.predict_proba(X_pred_imaged)
return y_pred_pro[0].tolist()
def get_text_prediction(imgpath):
result = reader.readtext(imgpath,paragraph="False")
text = []
for i in result:
text.append(i[1])
text = " ".join(text)
st.write(text)
t_pred = get_inference(text)
t_pred_c = []
for c in t_pred:
for a in c.values():
if a not in ['NEGATIVE','POSITIVE']:
t_pred_c.append(a)
return t_pred_c[::-1]
def pred_label_mean(i_pred,t_pred):
ensemble_pro = [(g + h) / 2 for g, h in zip(i_pred, t_pred)]
return ensemble_pro
def get_inference(input_text):
return bert(input_text)
loaded_lgbm = joblib.load('lgbm_v (2).sav')
bert = pipeline("text-classification", return_all_scores=True)
reader = easyocr.Reader(['en'])
st.title('Hateful Memes Classification')
image_file = st.file_uploader("Upload Images", type=["png","jpg","jpeg"])
temp_file = NamedTemporaryFile(delete=False)
if image_file is not None:
# To View Uploaded Image
st.write('Meme Image:')
temp_file.write(image_file.getvalue())
imgu = load_img(temp_file.name)
st.image(imgu)
with st.spinner('Predicting Label..'):
i_pred = get_img_prediction(temp_file.name)
t_pred = get_text_prediction(temp_file.name)
y_pred_both = pred_label_mean(i_pred,t_pred)
y_pred = y_pred_both.index(max(y_pred_both))
st.write(np.round(np.array(y_pred_both),4))
if y_pred == 0:
st.success('Predicted Label: non-hateful meme')
if y_pred == 1:
st.success('Predicted Label: hateful meme')