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

# Carregar o modelo de detecção de rostos
mtcnn = MTCNN(keep_all=True)

def detect_faces(image):
    # Detectar rostos na imagem
    boxes, _ = mtcnn.detect(image)
    
    # Desenhar as caixas delimitadoras na imagem
    draw = ImageDraw.Draw(image)
    num_faces = 0

    if boxes is not None:
        for box in boxes:
            draw.rectangle(box.tolist(), outline="red", width=3)
            num_faces += 1

    return image, f"Número de rostos: {num_faces}"

# Criar a interface do Gradio
iface = gr.Interface(
    fn=detect_faces, 
    inputs=gr.Image(type="pil"), 
    outputs=[gr.Image(type="pil"), gr.Textbox()],
    title="Contador de Rostos em Imagens",
    description="Carregue uma imagem para contar o número de rostos detectados nela."
)

# Executar a aplicação
iface.launch()