Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import io
-from PIL import Image
+from PIL import Image,ImageDraw
 from transformers import AutoImageProcessor, AutoModelForObjectDetection
 import streamlit as st
 import torch
@@ -45,7 +45,7 @@ if uploaded_file is not None:
     image = Image.open(uploaded_file)
     st.image(image, caption="Uploaded Image.", use_column_width=True)
 submit = st.button("Detect Objects ")
-if submit:
+"""if submit:
     image_data=input_image_setup(uploaded_file)
     st.subheader("The response is..")
     #process with model
@@ -65,3 +65,26 @@ if submit:
             f"Detected {model.config.id2label[label.item()]} with confidence "
             f"{round(score.item(), 3)} at location {box}"
         )
+"""
+if submit:
+    image_data = input_image_setup(uploaded_file)
+    st.subheader("The response is..")
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    logits = outputs.logits
+    bboxes = outputs.pred_boxes
+    
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]
+
+    # Draw bounding boxes on the image
+    drawn_image = image.copy()
+    draw = ImageDraw.Draw(drawn_image)
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        box = [int(i) for i in box.tolist()]
+        draw.rectangle(box, outline="red", width=2)
+        label_text = f"{model.config.id2label[label.item()]} ({round(score.item(), 2)})"
+        draw.text((box[0], box[1]), label_text, fill="red")
+    
+    st.image(drawn_image, caption="Detected Objects", use_column_width=True)