Update app.py

app.py

@@ -2,17 +2,32 @@ import streamlit as st
 import torch
 from torchvision.models.detection import fasterrcnn_resnet50_fpn
 from torchvision.transforms import functional as F
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageFont
 import io
+import cv2
+import numpy as np
 
-# Set page config
-st.set_page_config(page_title="Object Detection App", layout="centered")
+# Set up UI
+st.set_page_config(page_title="🎯 AI Object Detection", layout="centered")
 
-# Title and description
-st.title("🎯 AI Object Detection App")
-st.markdown("Upload an image, and let AI detect objects with bounding boxes using a pretrained Faster R-CNN model.")
+# COCO class labels (80 classes)
+COCO_INSTANCE_CATEGORY_NAMES = [
+    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
+    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
+    'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
+    'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A',
+    'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
+    'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
+    'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', 'cup', 'fork',
+    'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli',
+    'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
+    'potted plant', 'bed', 'N/A', 'dining table', 'N/A', 'N/A', 'toilet', 'N/A',
+    'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
+    'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book', 'clock', 'vase',
+    'scissors', 'teddy bear', 'hair drier', 'toothbrush'
+]
 
-# Load model
+# Load detection model
 @st.cache_resource
 def load_model():
     model = fasterrcnn_resnet50_fpn(pretrained=True)
@@ -21,46 +36,63 @@ def load_model():
 
 model = load_model()
 
-# Upload image
-uploaded_file = st.file_uploader("📷 Upload Image", type=["jpg", "jpeg", "png"])
+# Sidebar
+st.sidebar.title("⚙️ Settings")
+conf_thresh = st.sidebar.slider("Confidence Threshold", 0.2, 1.0, 0.5, 0.05)
+theme = st.sidebar.radio("Theme", ["Light", "Dark"])
+source = st.sidebar.radio("Input Source", ["Upload Image", "Use Webcam"])
 
-# Confidence threshold slider
-conf_thresh = st.slider("🎚 Confidence Threshold", min_value=0.1, max_value=1.0, value=0.5, step=0.05)
+# Main Title
+st.title("🎯 Object Detection with AI")
+st.markdown("Upload an image or use webcam to detect objects in real time using **Faster R-CNN** with bounding boxes and labels.")
 
-if uploaded_file is not None:
-    image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Original Image", use_column_width=True)
+def detect_objects(image, threshold=0.5):
+    tensor = F.to_tensor(image).unsqueeze(0)
+    outputs = model(tensor)
+    boxes = outputs[0]['boxes']
+    labels = outputs[0]['labels']
+    scores = outputs[0]['scores']
+    
+    draw = ImageDraw.Draw(image)
 
-    # Convert image to tensor
-    image_tensor = F.to_tensor(image).unsqueeze(0)
+    for box, label, score in zip(boxes, labels, scores):
+        if score >= threshold:
+            box = box.tolist()
+            name = COCO_INSTANCE_CATEGORY_NAMES[label.item()]
+            draw.rectangle(box, outline="red", width=3)
+            draw.text((box[0], box[1] - 10), f"{name} {score:.2f}", fill="white" if theme == "Dark" else "black")
+    return image
 
-    # Run detection
-    with st.spinner("Detecting objects..."):
-        outputs = model(image_tensor)
-        boxes = outputs[0]["boxes"]
-        labels = outputs[0]["labels"]
-        scores = outputs[0]["scores"]
+if source == "Upload Image":
+    uploaded_file = st.file_uploader("📤 Upload Image", type=["jpg", "jpeg", "png"])
+    if uploaded_file:
+        image = Image.open(uploaded_file).convert("RGB")
+        st.image(image, caption="Original Image", use_column_width=True)
+        st.subheader("🧠 Detected Objects:")
+        detected = detect_objects(image.copy(), conf_thresh)
+        st.image(detected, caption="Detection Result", use_column_width=True)
 
-    # Filter boxes by confidence threshold
-    selected_indices = [i for i, score in enumerate(scores) if score >= conf_thresh]
-    draw = ImageDraw.Draw(image)
+        # Download button
+        buf = io.BytesIO()
+        detected.save(buf, format="PNG")
+        st.download_button("📥 Download Image", data=buf.getvalue(), file_name="detected.png", mime="image/png")
+
+else:
+    st.subheader("📸 Real-Time Webcam Detection")
+    run = st.checkbox("Turn On Webcam")
+
+    FRAME_WINDOW = st.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")
+    cap = cv2.VideoCapture(0)
 
-    st.image(image, caption="🧠 Detected Image", use_column_width=True)
+    while run:
+        ret, frame = cap.read()
+        if not ret:
+            st.write("Failed to get webcam feed.")
+            break
 
-    # 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"
-    )
+        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        pil_img = Image.fromarray(img)
+        result = detect_objects(pil_img.copy(), conf_thresh)
+        FRAME_WINDOW.image(result)
+    cap.release()