Create CrowdDetection.py

CrowdDetection.py

@@ -0,0 +1,65 @@
+import gradio as gr
+import torch
+import cv2
+from ultralytics import YOLO
+
+# Crowd Detection Class
+class CrowdDetection:
+    def __init__(self, yolo_model_path="yolov8n.pt", crowd_threshold=10):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = YOLO(yolo_model_path).to(self.device)  # Move to GPU
+        self.model.nms = 0.5  # Set NMS threshold
+        self.crowd_threshold = crowd_threshold
+
+    def detect_crowd(self, video_path):
+        cap = cv2.VideoCapture(video_path)
+        output_path = "output_crowd.mp4"
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        out = cv2.VideoWriter(output_path, fourcc, int(cap.get(cv2.CAP_PROP_FPS)),
+                              (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))
+
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+
+            results = self.model(frame)
+            person_count = 0
+
+            for result in results:
+                boxes = result.boxes.xyxy.cpu().numpy()
+                classes = result.boxes.cls.cpu().numpy()
+
+                for box, cls in zip(boxes, classes):
+                    if int(cls) == 0:  # YOLO class ID 0 = "person"
+                        person_count += 1
+                        x1, y1, x2, y2 = map(int, box)
+                        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                        cv2.putText(frame, "Person", (x1, y1 - 10),
+                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+
+            alert_text = "Crowd Alert!" if person_count > self.crowd_threshold else f"People: {person_count}"
+            cv2.putText(frame, alert_text, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
+                        (0, 0, 255) if person_count > self.crowd_threshold else (0, 255, 0), 2)
+
+            out.write(frame)
+
+        cap.release()
+        out.release()
+        return output_path
+
+# Gradio Function
+def process_video(video):
+    detector = CrowdDetection()
+    return detector.detect_crowd(video)
+
+# Gradio Interface
+interface = gr.Interface(
+    fn=process_video,
+    inputs=gr.Video(label="Upload Video"),
+    outputs=gr.Video(label="Processed Video"),
+    title="Crowd Detection using YOLOv8"
+)
+
+if __name__ == "__main__":
+    interface.launch()