Update app.py

app.py

@@ -2,6 +2,9 @@ import streamlit as st
 import tensorflow as tf
 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_path):
@@ -22,6 +25,27 @@ if file_name is not None:
 
     input_image,copy_image = process_img(file_name)
     col1.image(input_image, use_column_width=True)
+    
+    bbox,confidance,classes,nc = model(input_img)
+    bbox,confidance , classes , nc  = bbox[0].numpy(),confidance[0].numpy(),classes[0].numpy(),nc[0].numpy()
+    
+    for i in range(nc):
+    if confidance[i] >= threshold:
+        x1,y1,x2,y2 = bbox[i]*640
+        class_name = labels[int(classes[i])]
+        if class_name =="Header":
+            color = (0,0,255) #Blue color
+            cv2.rectangle(copy_img, (int(x1), int(y1)), (int(x2), int(y2)),color, 2)
+
+        if class_name =="Column":
+            color = (0,255,0) #Green color
+            cv2.rectangle(copy_img, (int(x1), int(y1)), (int(x2), int(y2)),color, 2)
+
+        if class_name =="Table":
+            color = (255,0,0) #Red color
+            cv2.rectangle(copy_img, (int(x1), int(y1)), (int(x2), int(y2)),color, 2)
+
+    col2.image(copy_img, use_column_width=True)