import gradio as gr
import torch
import cv2
from ultralytics import YOLO
import spaces 
@spaces.GPU  


class CrowdDetection:
    def __init__(self, yolo_model_path="yolov8n.pt", crowd_threshold=10):
        # Determine the best available device
        if torch.cuda.is_available():
            self.device = torch.device("cuda")
            print(f"Using CUDA Device: {torch.cuda.get_device_name(0)}")
        else:
            self.device = torch.device("cpu")
            print("Using CPU as no CUDA device is available")
        
        try:
            # Load the YOLO model on the selected device
            self.model = YOLO(yolo_model_path).to(self.device)
        except Exception as e:
            print(f"Error loading YOLO model: {e}")
            self.model = None
        
        self.crowd_threshold = crowd_threshold

    def detect_crowd(self, video_path):
        if self.model is None:
            raise ValueError("YOLO model failed to load")
        
        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
            
            # Perform detection
            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

def process_video(video):
    try:
        detector = CrowdDetection()
        output_video = detector.detect_crowd(video)
        return output_video
    except Exception as e:
        print(f"Video processing error: {e}")
        return None

# Gradio Interface for Hugging Face Spaces
interface = gr.Interface(
    fn=process_video,
    inputs=gr.Video(label="Upload Video"),
    outputs=gr.Video(label="Processed Video"),
    title="Crowd Detection using YOLOv8"
)

# Remove share=True for Hugging Face Spaces
if __name__ == "__main__":
    interface.launch()