Update app.py

app.py

@@ -3,50 +3,48 @@ from PIL import Image
 import numpy as np
 from ultralytics import YOLO
 
-# Streamlit page config
 st.set_page_config(page_title="Suspicious Activity Detection", layout="centered")
 
 # Load the YOLOv11 model
 @st.cache_resource
 def load_model():
-    return YOLO("yolo11l.pt")  # Ensure the file is in your app directory
+    return YOLO("yolo11l.pt")  # Ensure your model is uploaded in the root of the Hugging Face Space
 
 model = load_model()
 
-# ------------------ Action Classification Logic ------------------
+# ------------------ Improved Classification Logic ------------------
 def classify_action(detections):
     action_scores = {'Stealing': 0, 'Sneaking': 0, 'Peaking': 0, 'Normal': 0}
     objects = [d[0] for d in detections]
     confidences = [d[1] for d in detections]
 
-    num_persons = objects.count('person')
-    total_objects = len(objects)
-    avg_conf = np.mean(confidences) if confidences else 0.0
-
-    has_valuable = any(obj in objects for obj in ['backpack', 'handbag', 'laptop', 'cash', 'microwave', 'refrigerator'])
-
-    if num_persons >= 1:
-        # Stealing
-        if has_valuable:
+    has_person = 'person' in objects
+    person_count = objects.count('person')
+    unique_objs = set(objects)
+    has_bag = any(obj in objects for obj in ['backpack', 'handbag'])
+    stealthy_conf = all(conf < 0.55 for conf in confidences) if confidences else True
+    low_activity = len(objects) <= 2
+
+    # --- Classify based on heuristic rules ---
+    if has_person:
+        # Stealing: person with object like bag/fridge and high confidence
+        if has_bag or any(obj in objects for obj in ['refrigerator', 'microwave', 'laptop']):
             action_scores['Stealing'] += 0.7
-            if 'handbag' in objects or 'laptop' in objects:
-                action_scores['Stealing'] += 0.2
 
-        # Sneaking - enhanced logic
-        if (total_objects <= 2 and avg_conf < 0.5 and not has_valuable):
+        # Sneaking: low confidence + very few objects + only person
+        if stealthy_conf and person_count == 1 and len(unique_objs) == 1:
             action_scores['Sneaking'] += 0.9
 
-        # Peaking
-        elif (total_objects <= 3 and num_persons == 1 and avg_conf >= 0.5):
-            action_scores['Peaking'] += 0.7
+        # Peaking: low object count, mostly person, no interaction
+        elif person_count >= 1 and low_activity:
+            action_scores['Peaking'] += 0.6
 
-        # Normal fallback
-        else:
-            action_scores['Normal'] += 0.6
+        # Otherwise: Normal
+        if action_scores['Sneaking'] == 0 and action_scores['Peaking'] == 0 and action_scores['Stealing'] == 0:
+            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:
@@ -54,24 +52,21 @@ def classify_action(detections):
 
     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("Detect if someone is **Stealing**, **Sneaking**, **Peaking**, or just acting **Normal**.")
+st.title("🕵️ Suspicious Activity Detection")
+st.markdown("Detect if someone is **Stealing**, **Sneaking**, **Peaking**, or just acting **Normal** based on image input.")
 
 uploaded_file = st.file_uploader("📸 Upload an image", type=["jpg", "jpeg", "png"])
 
@@ -82,7 +77,7 @@ if uploaded_file:
     temp_path = "/tmp/uploaded.jpg"
     image.save(temp_path)
 
-    with st.spinner("Detecting activity..."):
+    with st.spinner("🔎 Detecting activity..."):
         detected_image, action_scores = detect_action(temp_path)
 
     st.image(detected_image, caption="Detection Results", use_column_width=True)
@@ -92,4 +87,4 @@ if uploaded_file:
         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)")
+    st.success(f"🎯 **Predicted Action:** {top_action[0]} ({top_action[1]:.2%} confidence)")