Update app.py

app.py

@@ -2,16 +2,21 @@ import streamlit as st
 from transformers import DetrForObjectDetection, DetrImageProcessor
 from PIL import Image
 import torch
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
 
-# Load model and processor
+# Load the model and processor
 model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
 processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 
 st.title("Context-Aware Object Detection")
+st.write("Upload an image to detect objects with contextual awareness.")
 
 # Upload an image
-uploaded_file = st.file_uploader("Choose an image...", type="jpg")
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
 if uploaded_file is not None:
+    # Open the uploaded image
     image = Image.open(uploaded_file)
     st.image(image, caption="Uploaded Image", use_column_width=True)
 
@@ -19,15 +24,30 @@ if uploaded_file is not None:
     inputs = processor(images=image, return_tensors="pt")
     outputs = model(**inputs)
 
-    # Extract and display bounding boxes and labels
-    logits = outputs.logits.softmax(-1)
-    boxes = outputs.pred_boxes
+    # Get logits and bounding boxes
+    logits = outputs.logits.softmax(-1)[0]
+    boxes = outputs.pred_boxes[0]
 
-    # Define a confidence threshold
+    # Set a confidence threshold for displaying boxes
     threshold = 0.9
-    for logit, box in zip(logits[0], boxes[0]):
-        score, label = logit.max(0)
+    labels = processor.tokenizer.convert_ids_to_tokens(logits.argmax(-1))
+    scores = logits.max(-1).values
+
+    # Display the image with bounding boxes
+    fig, ax = plt.subplots(1)
+    ax.imshow(image)
+
+    # Plot each detected object if it meets the confidence threshold
+    for score, label, box in zip(scores, labels, boxes):
         if score > threshold:
-            st.write(f"Detected object with confidence {score:.2f}")
+            # Convert bounding box coordinates to absolute pixel values
+            x, y, w, h = box * torch.tensor([image.width, image.height, image.width, image.height])
+            x0, y0 = x - w / 2, y - h / 2
+
+            # Draw the bounding box
+            rect = patches.Rectangle((x0, y0), w, h, linewidth=2, edgecolor='r', facecolor='none')
+            ax.add_patch(rect)
+            ax.text(x0, y0, f"{label}: {score:.2f}", color='red', fontsize=8, weight='bold')
+
+    st.pyplot(fig)
 
-    st.write("Detection complete!")