FlowchartOCR / app.py
srijaydeshpande's picture
Update app.py
cf8530c verified
import pandas as pd
import numpy as np
import streamlit as st
import easyocr
import PIL
from PIL import Image, ImageDraw
def rectangle(image, result):
# https://www.blog.pythonlibrary.org/2021/02/23/drawing-shapes-on-images-with-python-and-pillow/
""" draw rectangles on image based on predicted coordinates"""
draw = ImageDraw.Draw(image)
for res in result:
top_left = tuple(res[0][0]) # top left coordinates as tuple
bottom_right = tuple(res[0][2]) # bottom right coordinates as tuple
draw.rectangle((top_left, bottom_right), outline="blue", width=2)
#display image on streamlit
st.image(image)
# main title
st.title("Recognize text and locations from flowchart")
# # subtitle
# st.markdown("## FlowchartOCR")
# upload image file
file = st.file_uploader(label = "Upload Image", type=['png', 'jpg', 'jpeg'])
#read the csv file and display the dataframe
if file is not None:
image = Image.open(file) # read image with PIL library
st.image(image) #display
# it will only detect the English and Turkish part of the image as text
reader = easyocr.Reader(['en'], gpu=False)
result = reader.readtext(np.array(image)) # turn image to numpy array
# collect the results in the dictionary:
textdic_easyocr = {}
for idx in range(len(result)):
pred_coor = result[idx][0]
x_coords = [point[0] for point in pred_coor]
y_coords = [point[1] for point in pred_coor]
center_x = sum(x_coords) / len(x_coords)
center_y = sum(y_coords) / len(y_coords)
pred_text = result[idx][1]
pred_confidence = result[idx][2]
if(pred_confidence>0.3):
textdic_easyocr[pred_text] = {}
textdic_easyocr[pred_text]['location'] = (center_x, center_y)
textdic_easyocr[pred_text]['pred_confidence'] = pred_confidence
# create a data frame which shows the predicted text and prediction confidence
df = pd.DataFrame.from_dict(textdic_easyocr).T
st.table(df)
# get boxes on the image
rectangle(image, result)
st.spinner(text="In progress...")
else:
st.write("Upload your image")