Update app.py

app.py

@@ -321,107 +321,132 @@ class IntrusionDetection:
             raise ValueError(f"Error in detect_intrusion: {str(e)}")
 
 
-class LoiteringDetection:
+import cv2
+import numpy as np
+from ultralytics import YOLO
+from shapely.geometry import Point, Polygon
+import time
+
+class LoiteringDetector:
+    def __init__(self, model_path='loitering_model.pt'):
+        self.model = YOLO(model_path)
     
-    def __init__(self, model_path="yolov8n.pt", loitering_threshold=10, conf_threshold=0.5):
-        self.model_path = model_path
-        self.loitering_threshold = loitering_threshold
-        self.conf_threshold = conf_threshold
-        self.entry_time = {}
-        self.area = [(153, 850), (139, 535), (239, 497), (291, 857)]
     @spaces.GPU
+    def loitering_detect(self, video_path, area):
+        # Create polygon zone
+        time_threshold = 7
+        detection_threshold = 0.6
+        zone_points = None
+        if area == '131':
+            zone_points = [(842//1.5, 514//1.7), (686//1.5, 290//1.7), (775//1.5, 279//1.7), (961//1.5, 488//1.7)]
+        elif area == '145':
+            zone_points = [(153, 850), (139, 535), (239, 497), (291, 857)]
+        zone = Polygon(zone_points)
+
+        # Open video
+        cap = cv2.VideoCapture(video_path)
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) * 0.5)
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) * 0.5)
+        fps = int(cap.get(cv2.CAP_PROP_FPS))
+
+        # Create video writer
+        output_path = os.path.join(tempfile.gettempdir(), "loitering_video.mp4")
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+        
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+
+            frame = cv2.resize(frame, (width, height))
+            # Perform object detection and tracking
+            results = self.model.track(frame, persist=True, classes=[0], conf=detection_threshold)  # 0 is the class ID for person
+
+            # List to store time information for display
+            time_display = []
+
+            if results[0].boxes.id is not None:
+                boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)
+                ids = results[0].boxes.id.cpu().numpy().astype(int)
 
-    def load_model(self):
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model = YOLO(self.model_path).to(device)
-        return model
+                for box, id in zip(boxes, ids):
+                    x1, y1, x2, y2 = box
+                    center = Point((x1 + x2) / 2, (y1 + y2) / 2)
 
-    def calculate_center(self, box):
-        x1, y1, x2, y2 = box
-        return int((x1 + x2) / 2), int((y1 + y2) / 2)
+                    if id not in person_info:
+                        person_info[id] = {'in_zone': False, 'start_time': None, 'duration': 0}
 
-    def track_time(self, id, frame_duration):
-        if id not in self.entry_time:
-            self.entry_time[id] = {'duration': 0, 'loitering': False}
-        else:
-            self.entry_time[id]['duration'] += frame_duration
-        if self.entry_time[id]['duration'] > self.loitering_threshold:
-            self.entry_time[id]['loitering'] = True
+                    if zone.contains(center):
+                        if not person_info[id]['in_zone']:
+                            person_info[id]['in_zone'] = True
+                            person_info[id]['start_time'] = time.time()
 
-    def loitering_detect(self, video_path): #edited from detect_loitering
-        try:
-            model = self.load_model()
-            cap = cv2.VideoCapture(video_path)
-            if not cap.isOpened():
-                raise ValueError(f"❌ Failed to open video: {video_path}")
+                        person_info[id]['duration'] = time.time() - person_info[id]['start_time']
 
-            fps = int(cap.get(cv2.CAP_PROP_FPS)) or 30
-            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            output_path = "output_loitering.mp4"
-            out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
-            frame_duration = 1 / fps
-            frame_count = 0
+                        if person_info[id]['duration'] > time_threshold:
+                            color = (0, 0, 255)  # Red for loitering
+                        else:
+                            color = (0, 255, 0)  # Green for in zone
 
-            while cap.isOpened():
-                ret, frame = cap.read()
-                if not ret:
-                    break
-                frame_count += 1
+                        time_display.append(f"ID: {id}, Time: {person_info[id]['duration']:.2f}s")
+                    else:
+                        person_info[id]['in_zone'] = False
+                        person_info[id]['start_time'] = None
+                        person_info[id]['duration'] = 0
+                        color = (255, 0, 0)  # Blue for outside zone
 
-                results = model.track(frame, conf=self.conf_threshold, iou=0.1, classes=[0], persist=True)
-                boxes = results[0].boxes.xyxy.cpu().tolist()
-                ids = results[0].boxes.id.cpu().tolist()
+                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+                    cv2.putText(frame, f"ID: {id}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
 
-                ids_in_area = []
-                for box, id in zip(boxes, ids):
-                    center = self.calculate_center(box)
-                    if cv2.pointPolygonTest(np.array(self.area, np.int32), center, False) >= 0:
-                        ids_in_area.append(id)
-                        self.track_time(id, frame_duration)
-
-                for id in ids_in_area:
-                    color = (0, 0, 255) if self.entry_time.get(id, {}).get('loitering', False) else (0, 255, 0)
-                    cv2.putText(frame, f"ID {id}, Time: {self.entry_time[id]['duration']:.1f}s", (15, 30 + id * 30),
-                                cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
-                
-                pts = np.array(self.area, np.int32).reshape((-1, 1, 2))
-                color = (0, 0, 255) if any(self.entry_time.get(id, {}).get('loitering', False) for id in ids_in_area) else (152, 251, 152)
-                cv2.polylines(frame, [pts], isClosed=True, color=color, thickness=3)
-                out.write(frame)
+            # Draw polygon zone
+            cv2.polylines(frame, [np.array(zone_points, np.int32)], True, (255, 255, 0), 2)
 
-            cap.release()
-            out.release()
-            if frame_count == 0 or not os.path.exists(output_path):
-                raise ValueError("❌ Processing failed: No frames processed or output not created")
-            return output_path
-        except Exception as e:
-            raise ValueError(f"Error in detect_loitering: {str(e)}")
+            # Display time information in top left
+            for i, text in enumerate(time_display):
+                cv2.putText(frame, text, (10, 30 + i * 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
+
+            out.write(frame)
+
+        cap.release()
+        out.release()
 
-# Unified processing function with status output
-def process_video(feature, video):
+        return output_path
+
+
+        def process_video(feature, video, area=None):
     detectors = {
         "Crowd Detection": CrowdDetection,
         "People Tracking": PeopleTracking,
         "Fall Detection": FallDetection,
         "Fight Detection": FightDetection,
-        "Intrusion Detection" : IntrusionDetection,
-        "Loitering Detection" : LoiteringDetection
+        "Intrusion Detection": IntrusionDetection,
+        "Loitering Detection": LoiteringDetection
     }
+    
     try:
         detector = detectors[feature]()
         method_name = feature.lower().replace(" ", "_").replace("detection", "detect")  # Ensures correct method name
-        output_path = getattr(detector, method_name)(video)
+        
+        if feature == "Loitering Detection":
+            output_path = detector.detect_loitering(video, area)  # Pass area if required
+        else:
+            output_path = getattr(detector, method_name)(video)
+        
         return f"{feature} completed successfully", output_path
     except Exception as e:
         return f"Error: {str(e)}", None
 
-# Gradio Interface with dual outputs
+# Gradio Interface with additional input for Loitering Detection
 interface = gr.Interface(
     fn=process_video,
     inputs=[
-        gr.Dropdown(choices=["Crowd Detection", "People Tracking", "Fall Detection", "Fight Detection", "Intrusion Detection", "Loitering Detection"], label="Select Feature"),
-        gr.Video(label="Upload Video")
+        gr.Dropdown(choices=[
+            "Crowd Detection", "People Tracking", "Fall Detection", 
+            "Fight Detection", "Intrusion Detection", "Loitering Detection"
+        ], label="Select Feature"),
+        gr.Video(label="Upload Video"),
+        gr.Textbox(label="Loitering Area (131 or 145)")  # Can be replaced with a drawing tool
     ],
     outputs=[
         gr.Textbox(label="Status"),
@@ -432,4 +457,4 @@ interface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    interface.launch(debug=True)
+    interface.launch(debug=True)