Update app.py

app.py

@@ -1,44 +1,99 @@
 import streamlit as st
 from PIL import Image
 import numpy as np
-import cv2
 from ultralytics import YOLO
 
-# Title
-st.title("Suspicious Activity Detection")
-st.markdown("Upload an image to detect activity (Normal, Peaking, Sneaking, Stealing).")
+# Set Streamlit page config
+st.set_page_config(page_title="Suspicious Activity Detection", layout="centered")
 
-# Load model
+# Load the YOLOv11 model
 @st.cache_resource
 def load_model():
-    model = YOLO("yolo11l.pt")  # Make sure this file is uploaded to your Space
-    return model
+    return YOLO("yolo11l.pt")  # Make sure yolo11l.pt is uploaded in your Space
 
 model = load_model()
 
-# File uploader
-uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "jpeg", "png"])
+# ------------------ Modified Action Classification ------------------
+def classify_action(detections):
+    """
+    Improved classification to distinguish Sneaking from Peaking.
+    """
+    action_scores = {
+        'Stealing': 0,
+        'Sneaking': 0,
+        'Peaking': 0,
+        'Normal': 0
+    }
+
+    objects = [d[0] for d in detections]
+    confidences = [d[1] for d in detections]
+
+    has_person = 'person' in objects
+    avg_conf = np.mean(confidences) if confidences else 0
+    num_objects = len(set(objects))
+
+    # Stealing Logic
+    if has_person and ('handbag' in objects or 'backpack' in objects or 'suitcase' in objects):
+        action_scores['Stealing'] += 0.6
+    if has_person and ('refrigerator' in objects or 'microwave' in objects or 'tv' in objects):
+        action_scores['Stealing'] += 0.4
+
+    # Sneaking Logic: person present, low confidence, few objects
+    if has_person and avg_conf < 0.5 and num_objects <= 2:
+        action_scores['Sneaking'] += 0.9
+
+    # Peaking Logic: moderate confidence, minimal visible objects
+    if has_person and 0.5 <= avg_conf <= 0.75 and num_objects <= 1:
+        action_scores['Peaking'] += 0.7
+
+    # Normal if nothing else fits
+    if all(score == 0 for score in action_scores.values()):
+        action_scores['Normal'] = 1.0
+
+    # Normalize
+    total = sum(action_scores.values())
+    if total > 0:
+        for key in action_scores:
+            action_scores[key] /= total
+
+    return action_scores
+
+# ------------------ Detection Function ------------------
+def detect_action(image_path):
+    results = model.predict(source=image_path, conf=0.4, iou=0.5, save=False, verbose=False)
+    result = results[0]
+
+    detections = [
+        (model.names[int(cls)], float(conf))
+        for cls, conf in zip(result.boxes.cls, result.boxes.conf)
+    ]
+
+    annotated_image = result.plot()
+    action_scores = classify_action(detections)
+
+    return annotated_image, action_scores
+
+# ------------------ Streamlit UI ------------------
+st.title("🔍 Suspicious Activity Detection")
+st.markdown("Upload an image to detect if someone is **Stealing**, **Sneaking**, **Peaking**, or just acting **Normal**.")
+
+uploaded_file = st.file_uploader("📤 Upload an image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file:
     image = Image.open(uploaded_file).convert("RGB")
     st.image(image, caption="Uploaded Image", use_column_width=True)
 
-    # Convert to numpy array for OpenCV/YOLO
-    img_array = np.array(image)
-
-    # Run prediction
-    st.info("Running detection...")
-    results = model.predict(source=img_array, conf=0.25, classes=None)
-
-    for r in results:
-        # Show annotated image
-        annotated_img = r.plot()
-        st.image(annotated_img, caption="Detection Result", use_column_width=True)
-
-        # Show detected objects with confidence
-        st.subheader("Detections:")
-        for box in r.boxes:
-            conf = float(box.conf[0])
-            cls_id = int(box.cls[0])
-            cls_name = model.names[cls_id]
-            st.write(f"- **{cls_name}** (Confidence: {conf:.2f})")
+    temp_path = "/tmp/uploaded.jpg"
+    image.save(temp_path)
+
+    with st.spinner("🕵️ Detecting suspicious activity..."):
+        detected_image, action_scores = detect_action(temp_path)
+
+    st.image(detected_image, caption="🔍 Detected Image", use_column_width=True)
+
+    st.subheader("📊 Action Confidence Scores")
+    for action, score in action_scores.items():
+        st.write(f"**{action}**: {score:.2%}")
+
+    most_likely_action = max(action_scores.items(), key=lambda x: x[1])
+    st.success(f"🎯 **Predicted Action**: {most_likely_action[0]} ({most_likely_action[1]:.2%} confidence)")