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)