Update app.py

app.py

@@ -1,28 +1,41 @@
 # Object Detection
+import streamlit as st
 from huggingface_hub import hf_hub_download
 from transformers import AutoImageProcessor, TableTransformerForObjectDetection
 import torch
 from PIL import Image
 
-file_path = hf_hub_download(repo_id="nielsr/example-pdf", repo_type="dataset", filename="example_pdf.png")
-image = Image.open(file_path).convert("RGB")
+# Model and Image Processor Loading (ideally at the app start)
+@st.cache_resource
+def load_assets():
+    file_path = hf_hub_download(repo_id="nielsr/example-pdf", repo_type="dataset", filename="example_pdf.png")
+    image_processor = AutoImageProcessor.from_pretrained("microsoft/table-transformer-detection")
+    model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-detection")
+    return file_path, image_processor, model
 
-image_processor = AutoImageProcessor.from_pretrained("microsoft/table-transformer-detection")
-model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-detection")
+file_path, image_processor, model = load_assets()
 
-inputs = image_processor(images=image, return_tensors="pt")
-outputs = model(**inputs)
+# App Title
+st.title("Table Detection in Images")
 
-# convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
-target_sizes = torch.tensor([image.size[::-1]])
-results = image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[
-    0
-]
+# Image Upload
+uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
 
-for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-    box = [round(i, 2) for i in box.tolist()]
-    print(
-        f"Detected {model.config.id2label[label.item()]} with confidence "
-        f"{round(score.item(), 3)} at location {box}"
-    )
-Detected table with confidence 1.0 at location [202.1, 210.59, 1119.22, 385.09]
+# Process Image and Display Results
+if uploaded_file:
+    image = Image.open(uploaded_file).convert("RGB")
+
+    inputs = image_processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]
+
+    st.image(image)  # Display the uploaded image
+
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        box = [round(i, 2) for i in box.tolist()]
+        st.write(
+            f"Detected {model.config.id2label[label.item()]} with confidence "
+            f"{round(score.item(), 3)} at location {box}"
+        )