Update app.py

app.py

@@ -215,72 +215,134 @@ class FallDetection:
             raise ValueError(f"Error in fall_detection: {str(e)}")
 
 
+import os
+import cv2
+import time
+import torch
+import numpy as np
+from ultralytics import YOLO
+
 class FightDetection:
     def __init__(self, yolo_model_path="yolov8n-pose.pt"):
         self.model_path = yolo_model_path
+
     
     @spaces.GPU
     def fight_detect(self, video_path):
         try:
             device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            
+            # Load YOLO Pose model
             if not os.path.exists(self.model_path):
                 model = YOLO("yolov8n-pose.pt")
                 model.save(self.model_path)
             else:
                 model = YOLO(self.model_path)
+            
             model.to(device)
 
             cap = cv2.VideoCapture(video_path)
             if not cap.isOpened():
                 raise ValueError(f"❌ Failed to open video: {video_path}")
 
-            fps = int(cap.get(cv2.CAP_PROP_FPS))
+            fps = int(cap.get(cv2.CAP_PROP_FPS)) // 2  # Slow down FPS for better tracking
             width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) * 0.5)
             height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) * 0.5)
             output_path = "output_fight.mp4"
+
             out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
             if not out.isOpened():
                 cap.release()
                 raise ValueError(f"❌ Failed to initialize video writer")
 
+            # Fight detection parameters
+            FIGHT_THRESHOLD = 2.0
+            PROXIMITY_THRESHOLD = 100
+            frame_skip = 2
+            frame_count = 0
+            person_movements = {}
+
             while cap.isOpened():
                 ret, frame = cap.read()
                 if not ret:
-                    break
-                
+                    break  # End of video
+
+                frame_count += 1
+                if frame_count % frame_skip != 0:
+                    continue  # Skip frames for performance
+
                 frame = cv2.resize(frame, (width, height))
                 results = model.track(frame, persist=True)
-                fight_detected = False
-                person_count = 0
+
+                current_time = time.time()
+                persons = []
 
                 for result in results:
                     keypoints = result.keypoints.xy.cpu().numpy() if result.keypoints else []
-                    boxes = result.boxes.xyxy.cpu().numpy() if result.boxes else []
                     classes = result.boxes.cls.cpu().numpy() if result.boxes else []
+                    ids = result.boxes.id.cpu().numpy() if result.boxes.id is not None else []
 
-                    for box, kp, cls in zip(boxes, keypoints, classes):
-                        if int(cls) == 0:
-                            person_count += 1
-                            x1, y1, x2, y2 = map(int, box)
-                            if len(kp) > 7 and (kp[5][1] < y1 + (y2 - y1) * 0.3 or kp[7][1] < y1 + (y2 - y1) * 0.3):
+                    for i, (kp, cls) in enumerate(zip(keypoints, classes)):
+                        if int(cls) == 0:  # Person class
+                            person_id = int(ids[i]) if len(ids) > i else f"{int(kp[0][0])}-{int(kp[0][1])}"
+                            persons.append((person_id, kp))
+
+                            if person_id not in person_movements:
+                                person_movements[person_id] = []
+
+                            person_movements[person_id].append((current_time, kp))
+
+                            # Draw keypoints
+                            for point in kp:
+                                x, y = int(point[0]), int(point[1])
+                                cv2.circle(frame, (x, y), 5, (255, 255, 0), -1)
+
+                # Check for fights
+                fight_detected = False
+                for i in range(len(persons)):
+                    for j in range(i + 1, len(persons)):
+                        person1, kp1 = persons[i]
+                        person2, kp2 = persons[j]
+
+                        distance = np.linalg.norm(kp1[0] - kp2[0])
+                        if distance > PROXIMITY_THRESHOLD:
+                            continue  # Ignore if too far apart
+
+                        if len(person_movements[person1]) > 1 and len(person_movements[person2]) > 1:
+                            hands1 = np.mean(kp1[[7, 8]], axis=0)
+                            hands2 = np.mean(kp2[[7, 8]], axis=0)
+
+                            prev_hands1 = person_movements[person1][-2][1][[7, 8]].mean(axis=0)
+                            prev_hands2 = person_movements[person2][-2][1][[7, 8]].mean(axis=0)
+
+                            speed1 = np.linalg.norm(hands1 - prev_hands1)
+                            speed2 = np.linalg.norm(hands2 - prev_hands2)
+
+                            if speed1 > FIGHT_THRESHOLD and speed2 > FIGHT_THRESHOLD:
                                 fight_detected = True
-                            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255) if fight_detected else (0, 255, 0), 2)
-                            label = "FIGHT DETECTED" if fight_detected else "Person"
-                            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, 
-                                        (0, 0, 255) if fight_detected else (0, 255, 0), 2)
+                                x1, y1 = int(kp1[0][0]), int(kp1[0][1])
+                                x2, y2 = int(kp2[0][0]), int(kp2[0][1])
+                                cv2.line(frame, (x1, y1), (x2, y2), (0, 0, 255), 3)
+                                cv2.putText(frame, "FIGHT DETECTED", (x1, y1 - 10), 
+                                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
 
-                if fight_detected and person_count > 1:
+                if fight_detected:
                     cv2.putText(frame, "FIGHT ALERT!", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+
                 out.write(frame)
 
             cap.release()
             out.release()
+
             if not os.path.exists(output_path):
                 raise ValueError("❌ Processing failed")
+
             return output_path
+
         except Exception as e:
             raise ValueError(f"Error in fight_detection: {str(e)}")
 
+
 class IntrusionDetection:
     def __init__(self, model_path="yolov8n.pt", max_intrusion_time=300, iou_threshold=0.5, conf_threshold=0.5):
         self.model_path = model_path