import gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
from supervision import Detections
from PIL import Image, ImageDraw

# Download & load the YOLOv8 face detection model
model_path = hf_hub_download(
    repo_id="arnabdhar/YOLOv8-Face-Detection",
    filename="model.pt"
)
model = YOLO(model_path)

def count_faces(image: Image.Image):
    """
    Detects faces in an image and returns the annotated image + face count.
    """
    # 1. Run detection
    results = model(image)[0]
    dets = Detections.from_ultralytics(results)

    # 2. Draw boxes
    annotated = image.copy()
    draw = ImageDraw.Draw(annotated)
    for x1, y1, x2, y2 in dets.xyxy:
        draw.rectangle([x1, y1, x2, y2], outline="red", width=2)

    # 3. Return
    return annotated, f"Faces detected: {len(dets.xyxy)}"

# Gradio interface
image_app = gr.Interface(
    fn=count_faces,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Image(type="pil", label="Annotated Image"),
        gr.Text(label="Face Count")
    ],
    title="Image Face Counter",
    description="Detect and count faces in a single image using YOLOv8."
)

if __name__ == "__main__":
    image_app.launch()