import gradio as gr
from PIL import Image, ImageDraw
from facenet_pytorch import MTCNN
import torch

# Load MTCNN model
device = 'cuda' if torch.cuda.is_available() else 'cpu'
mtcnn = MTCNN(keep_all=True, device=device)

# Face detection and density classification function
def detect_faces(image):
    if image is None:
        return None, "No image provided.", "No density calculated."
    
    # Detect faces
    boxes, _ = mtcnn.detect(image)
    face_count = 0 if boxes is None else len(boxes)

    # Classify density
    if face_count <= 10:
        density = "🟢 Sparse"
    elif face_count <= 50:
        density = "🟡 Medium"
    else:
        density = "🔴 Dense"

    # Annotate image
    annotated = image.copy()
    draw = ImageDraw.Draw(annotated)
    if boxes is not None:
        for box in boxes:
            draw.rectangle(box.tolist(), outline="red", width=3)

    # Return annotated image and stats
    return annotated, f"🧮 Face Count: {face_count}", f"📊 Crowd Density: {density}"

# Gradio UI
title = "🎯 Face Counter & Density Estimator"
description = """
Upload an image or use your webcam to detect faces and estimate crowd density.
"""

iface = gr.Interface(
    fn=detect_faces,
    inputs=gr.Image(sources=["upload", "webcam"], type="pil", label="Upload Image or use Webcam"),
    outputs=[
        gr.Image(type="pil", label="Detected Faces"),
        gr.Textbox(label="Face Count"),
        gr.Textbox(label="Crowd Density"),
    ],
    title=title,
    description=description,
    allow_flagging="never",
    theme="soft"
)

iface.launch()