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()