Update app.py

app.py

@@ -1,71 +1,89 @@
 import streamlit as st
-import torch
 from PIL import Image
-import numpy as np
-import cv2
-import tempfile
-import os
+from ultralytics import YOLO
+
+# Set up Streamlit page
+st.set_page_config(page_title="Suspicious Activity Detection", layout="centered")
 
 # Load YOLOv11 model
 @st.cache_resource
 def load_model():
-    return torch.hub.load('ultralytics/yolov5', 'custom', path='yolo11l.pt', force_reload=True)
+    return YOLO("yolo11l (1).pt")  # Ensure model filename matches
 
 model = load_model()
-model.conf = 0.25  # confidence threshold
 
-# Activity labels
-labels = ['Normal', 'Peaking', 'Sneaking', 'Stealing']
+# ------------------ Improved Action Classification Logic ------------------
+def classify_action(detections):
+    """
+    Classify activity based on object types, count, and confidence.
+    """
+    action_scores = {'Stealing': 0, 'Sneaking': 0, 'Peaking': 0, 'Normal': 0}
+    objects = [d[0] for d in detections]
+    confidences = [d[1] for d in detections]
 
-# Streamlit UI
-st.set_page_config(page_title="Suspicious Activity Detection", layout="centered")
-st.title("🚨 Suspicious Activity Detection with YOLOv11")
-st.markdown("Detect **Normal**, **Peaking**, **Sneaking**, and **Stealing** behaviors in real-time from uploaded images or videos.")
-
-# Upload section
-file = st.file_uploader("Upload an image or a video", type=['jpg', 'jpeg', 'png', 'mp4'])
-
-# Process predictions
-def display_predictions(results):
-    df = results.pandas().xyxy[0]
-    activity_counts = {label: 0 for label in labels}
-    for _, row in df.iterrows():
-        label = row['name']
-        if label in activity_counts:
-            activity_counts[label] += 1
-    return activity_counts
-
-# Handle image
-if file is not None:
-    if file.type.startswith("image"):
-        image = Image.open(file).convert('RGB')
-        st.image(image, caption="Uploaded Image", use_column_width=True)
-        results = model(image, size=640)
-        results.render()
-        st.image(Image.fromarray(results.ims[0]), caption="Detection Result", use_column_width=True)
-        counts = display_predictions(results)
-        st.markdown("### 🔍 Detected Activities")
-        for act, cnt in counts.items():
-            if cnt > 0:
-                st.success(f"**{act}**: {cnt}")
-
-    # Handle video
-    elif file.type.startswith("video"):
-        tfile = tempfile.NamedTemporaryFile(delete=False)
-        tfile.write(file.read())
-        video_path = tfile.name
-        cap = cv2.VideoCapture(video_path)
-        stframe = st.empty()
-        st.markdown("### 📹 Processing video...")
-
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-            results = model(frame)
-            results.render()
-            output_frame = results.ims[0]
-            stframe.image(output_frame, channels="BGR", use_column_width=True)
-        
-        cap.release()
-        os.remove(video_path)
+    has_person = 'person' in objects
+    has_bag = any(obj in objects for obj in ['handbag', 'backpack'])
+    few_objects = len(set(objects)) <= 2
+    mostly_person = objects.count('person') >= len(objects) * 0.6 if objects else False
+    max_conf = max(confidences) if confidences else 0.0
+
+    # Decision tree for classification
+    if has_person:
+        if has_bag and len(objects) >= 3:
+            action_scores['Stealing'] += 1.0
+        elif max_conf < 0.55 and few_objects:
+            action_scores['Sneaking'] += 1.0
+        elif mostly_person and few_objects and max_conf >= 0.55:
+            action_scores['Peaking'] += 1.0
+        else:
+            action_scores['Normal'] += 1.0
+    else:
+        action_scores['Normal'] += 1.0
+
+    # Normalize scores
+    total = sum(action_scores.values())
+    if total > 0:
+        for k in action_scores:
+            action_scores[k] /= total
+
+    return action_scores
+
+# ------------------ Detection Function ------------------
+def detect_action(image_path):
+    results = model.predict(source=image_path, conf=0.35, 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 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)
+
+    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="🔍 Detection Results", use_column_width=True)
+
+    st.subheader("📊 Action Confidence Scores")
+    for action, score in action_scores.items():
+        st.write(f"**{action}**: {score:.2%}")
+
+    top_action = max(action_scores.items(), key=lambda x: x[1])
+    st.success(f"🎯 **Predicted Action:** {top_action[0]} ({top_action[1]:.2%} confidence)")