Update app.py

app.py

@@ -1,89 +1,71 @@
 import streamlit as st
+import torch
 from PIL import Image
-from ultralytics import YOLO
-
-# Set up Streamlit page
-st.set_page_config(page_title="Suspicious Activity Detection", layout="centered")
+import numpy as np
+import cv2
+import tempfile
+import os
 
 # Load YOLOv11 model
 @st.cache_resource
 def load_model():
-    return YOLO("yolo11l.pt")  # Ensure model filename matches
+    return torch.hub.load('ultralytics/yolov5', 'custom', path='yolo11l.pt', force_reload=True)
 
 model = load_model()
+model.conf = 0.25  # confidence threshold
 
-# ------------------ 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]
-
-    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)
+# Activity labels
+labels = ['Normal', 'Peaking', 'Sneaking', 'Stealing']
 
-    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)")
+# 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)