Update app.py

app.py

@@ -1,92 +1,65 @@
 import streamlit as st
 from PIL import Image
 import torch
-from torchvision import transforms
-from torchvision.models.detection import fasterrcnn_resnet50_fpn
+import torchvision.transforms as T
 import torchvision
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
 
-# Load model
-model = fasterrcnn_resnet50_fpn(pretrained=True)
+# Load pre-trained object detection model
+model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
 model.eval()
 
-# Define class labels
+# COCO class labels
 COCO_INSTANCE_CATEGORY_NAMES = [
     '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
-    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
+    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', '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'
+    'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag',
+    'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite',
+    'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
+    'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana',
+    'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
+    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table',
+    'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
+    'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
+    'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
 ]
 
-def get_prediction(img, threshold):
-    transform = transforms.Compose([transforms.ToTensor()])
-    img = transform(img)
-    pred = model([img])
-    pred_classes = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in list(pred[0]['labels'].numpy())]
-    pred_boxes = [[(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]['boxes'].detach().numpy())]
-    pred_score = list(pred[0]['scores'].detach().numpy())
-    pred_t = [pred_score.index(x) for x in pred_score if x > threshold][-1]
-    boxes = pred_boxes[:pred_t+1]
-    classes = pred_classes[:pred_t+1]
-    return boxes, classes
+# Streamlit UI
+st.set_page_config(page_title="AI Object Detector", layout="centered")
+st.title("🎯 Object Detection with AI")
+st.markdown("Upload an image and let the AI detect objects with names!")
 
-# UI design
-st.set_page_config(page_title="AI Object Detector", layout="wide")
+uploaded_file = st.file_uploader("📷 Upload an image", type=["jpg", "png", "jpeg"])
 
-st.markdown("""
-    <style>
-        .main {
-            background-color: #f5f7fa;
-            padding: 20px;
-            border-radius: 10px;
-        }
-        h1 {
-            color: #2c3e50;
-        }
-        .stButton>button {
-            background-color: #008CBA;
-            color: white;
-            font-weight: bold;
-            border-radius: 8px;
-            padding: 10px 24px;
-        }
-    </style>
-""", unsafe_allow_html=True)
-
-st.title("🔍 AI Object Detection App")
-st.markdown("Upload an image and let the AI detect what's in it!")
-
-img_file = st.file_uploader("📸 Upload an Image", type=["jpg", "jpeg", "png"])
-
-confidence = st.slider("🎯 Confidence Threshold", 0.0, 1.0, 0.5)
-
-if img_file is not None:
-    image = Image.open(img_file).convert("RGB")
+if uploaded_file:
+    image = Image.open(uploaded_file).convert("RGB")
     st.image(image, caption="Uploaded Image", use_column_width=True)
 
-    boxes, classes = get_prediction(image, confidence)
+    # Transform image for model input
+    transform = T.Compose([T.ToTensor()])
+    image_tensor = transform(image).unsqueeze(0)
 
-    # Draw results
-    import matplotlib.pyplot as plt
-    import matplotlib.patches as patches
+    with st.spinner("🔍 Detecting objects..."):
+        predictions = model(image_tensor)[0]
 
-    fig, ax = plt.subplots(1, figsize=(12, 8))
+    # Draw bounding boxes
+    fig, ax = plt.subplots(1)
     ax.imshow(image)
-    for i in range(len(boxes)):
-        box = boxes[i]
-        label = classes[i]
-        rect = patches.Rectangle(box[0], box[1][0]-box[0][0], box[1][1]-box[0][1],
-                                 linewidth=2, edgecolor='blue', facecolor='none')
-        ax.add_patch(rect)
-        ax.text(box[0][0], box[0][1]-10, label, fontsize=12,
-                color='black', bbox=dict(facecolor='lightblue', edgecolor='blue', boxstyle='round,pad=0.5'))
+
+    threshold = 0.7  # confidence threshold
+    for idx in range(len(predictions["boxes"])):
+        score = predictions["scores"][idx].item()
+        if score > threshold:
+            box = predictions["boxes"][idx].detach().numpy()
+            label = COCO_INSTANCE_CATEGORY_NAMES[predictions["labels"][idx]]
+            x1, y1, x2, y2 = box
+            rect = patches.Rectangle((x1, y1), x2 - x1, y2 - y1,
+                                     linewidth=2, edgecolor='lime', facecolor='none')
+            ax.add_patch(rect)
+            ax.text(x1, y1 - 10, f"{label} ({score:.2f})", color='lime',
+                    fontsize=10, backgroundcolor='black')
+
     st.pyplot(fig)