TableExtract / app.py
905saini's picture
Update app.py
bae448a
import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import cv2
labels = ["Column","Header","Table"]
threshold = 0.75
model = tf.saved_model.load("best_saved_model") #Loading the saved model
def process_img(img):
img = img.resize((640,640)) # resize the image as reqired for the model input
img = np.array(img) # Convert Pil Object into numpy array
copy_img = img.copy() # Save a copy for backup
img = img/255 # Normalizing the pixel value
img = img.astype("float32") # Convert the format double to format float
img = np.expand_dims(img,axis=0) # exapanding dimension to add batch
return img,copy_img
st.title("Table Extract")
file_name = st.file_uploader("Upload a report image")
if file_name is not None:
col1, col2 = st.columns(2)
image = Image.open(file_name)
input_image, copy_img= process_img(image)
col1.header("Input Image")
col1.image(input_image, use_column_width=True)
bbox,confidance,classes,nc = model(input_image)
bbox,confidance , classes , nc = bbox[0].numpy(),confidance[0].numpy(),classes[0].numpy(),nc[0].numpy()
st.subheader("Detected Result")
table_count =0
header_count = 0
column_count = 0
for i in range(nc):
if confidance[i] >= threshold:
x1,y1,x2,y2 = bbox[i]*640
class_name = labels[int(classes[i])]
st.text(class_name+" : "+str(int(confidance[i]*100))+"%")
if class_name =="Header":
header_count+=1
color = (0,0,255) #Blue color
cv2.rectangle(copy_img, (int(x1), int(y1)), (int(x2), int(y2)),color, 2)
if class_name =="Column":
column_count+=1
color = (0,255,0) #Green color
cv2.rectangle(copy_img, (int(x1), int(y1)), (int(x2), int(y2)),color, 2)
if class_name =="Table":
table_count+=1
color = (255,0,0) #Red color
cv2.rectangle(copy_img, (int(x1), int(y1)), (int(x2), int(y2)),color, 2)
st.text("No of Table Detected : "+str(table_count))
st.text("No of Header Detected : "+str(header_count))
st.text("No of Column Detected : "+str(column_count))
col2.header("Output Result")
col2.image(copy_img, use_column_width=True)