Update app.py

app.py

@@ -1,49 +1,66 @@
-# app.py
 import streamlit as st
-from PIL import Image
 import torch
-import torchvision.transforms as T
-import requests
-from io import BytesIO
+from torchvision.models.detection import fasterrcnn_resnet50_fpn
+from torchvision.transforms import functional as F
+from PIL import Image, ImageDraw
+import io
 
-# Load a pre-trained model (Faster R-CNN)
-model = torch.hub.load('pytorch/vision:v0.10.0', 'fasterrcnn_resnet50_fpn', pretrained=True)
-model.eval()
+# Set page config
+st.set_page_config(page_title="Object Detection App", layout="centered")
 
-# COCO dataset labels
-LABELS_URL = 'https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt'
-labels = requests.get(LABELS_URL).text.split('\n')
-
-# Streamlit UI
-st.set_page_config(page_title="🎯 Object Detector", layout="centered")
+# Title and description
 st.title("🎯 AI Object Detection App")
-st.markdown("Detect objects in your images using a pre-trained AI model.")
+st.markdown("Upload an image, and let AI detect objects with bounding boxes using a pretrained Faster R-CNN model.")
+
+# Load model
+@st.cache_resource
+def load_model():
+    model = fasterrcnn_resnet50_fpn(pretrained=True)
+    model.eval()
+    return model
+
+model = load_model()
 
-uploaded_file = st.file_uploader("📤 Upload an Image", type=["jpg", "jpeg", "png"])
+# Upload image
+uploaded_file = st.file_uploader("📷 Upload Image", type=["jpg", "jpeg", "png"])
 
-if uploaded_file:
+# Confidence threshold slider
+conf_thresh = st.slider("🎚 Confidence Threshold", min_value=0.1, max_value=1.0, value=0.5, step=0.05)
+
+if uploaded_file is not None:
     image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-
-    # Transform image
-    transform = T.Compose([T.ToTensor()])
-    input_tensor = transform(image).unsqueeze(0)
-
-    with st.spinner('Detecting objects...'):
-        outputs = model(input_tensor)[0]
-
-    threshold = st.slider("Confidence Threshold", 0.0, 1.0, 0.5, 0.05)
-
-    draw = image.copy()
-    draw_edit = ImageDraw.Draw(draw)
-    
-    for idx, score in enumerate(outputs['scores']):
-        if score > threshold:
-            box = outputs['boxes'][idx].detach().numpy()
-            label_id = int(outputs['labels'][idx])
-            label = labels[label_id] if label_id < len(labels) else str(label_id)
-            draw_edit.rectangle(box, outline="red", width=3)
-            draw_edit.text((box[0], box[1]), f"{label}: {score:.2f}", fill="white")
-
-    st.image(draw, caption="Detected Objects", use_column_width=True)
-    st.success("Detection complete!")
+    st.image(image, caption="Original Image", use_column_width=True)
+
+    # Convert image to tensor
+    image_tensor = F.to_tensor(image).unsqueeze(0)
+
+    # Run detection
+    with st.spinner("Detecting objects..."):
+        outputs = model(image_tensor)
+        boxes = outputs[0]["boxes"]
+        labels = outputs[0]["labels"]
+        scores = outputs[0]["scores"]
+
+    # Filter boxes by confidence threshold
+    selected_indices = [i for i, score in enumerate(scores) if score >= conf_thresh]
+    draw = ImageDraw.Draw(image)
+
+    for i in selected_indices:
+        box = boxes[i].tolist()
+        label = labels[i].item()
+        score = scores[i].item()
+        draw.rectangle(box, outline="red", width=3)
+        draw.text((box[0], box[1]), f"{label}:{score:.2f}", fill="white")
+
+    st.image(image, caption="🧠 Detected Image", use_column_width=True)
+
+    # Download button
+    buf = io.BytesIO()
+    image.save(buf, format="PNG")
+    byte_im = buf.getvalue()
+    st.download_button(
+        label="📥 Download Detected Image",
+        data=byte_im,
+        file_name="detected.png",
+        mime="image/png"
+    )