Update app.py

app.py

@@ -1,10 +1,11 @@
 import streamlit as st
-from transformers import DetrForObjectDetection, DetrFeatureExtractor
+from transformers import DetrForObjectDetection, DetrImageProcessor
 from PIL import Image
+import torch
 
-# Load model and feature extractor
+# Load model and processor
 model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
-feature_extractor = DetrFeatureExtractor.from_pretrained("facebook/detr-resnet-50")
+processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 
 st.title("Context-Aware Object Detection")
 
@@ -14,8 +15,19 @@ if uploaded_file is not None:
     image = Image.open(uploaded_file)
     st.image(image, caption="Uploaded Image", use_column_width=True)
 
-    # Run detection
-    inputs = feature_extractor(images=image, return_tensors="pt")
+    # Preprocess the image and make predictions
+    inputs = processor(images=image, return_tensors="pt")
     outputs = model(**inputs)
-    # Display results here
-    st.write(outputs)
+
+    # Extract and display bounding boxes and labels
+    logits = outputs.logits.softmax(-1)
+    boxes = outputs.pred_boxes
+
+    # Define a confidence threshold
+    threshold = 0.9
+    for logit, box in zip(logits[0], boxes[0]):
+        score, label = logit.max(0)
+        if score > threshold:
+            st.write(f"Detected object with confidence {score:.2f}")
+
+    st.write("Detection complete!")