import gradio as gr
from transformers import DetrImageProcessor, DetrForObjectDetection
import torch
from PIL import Image, ImageDraw

# Carregar o processador e o modelo
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")

def detect_people(image):
    # Pré-processar a imagem
    inputs = processor(images=image, return_tensors="pt")

    # Realizar a inferência
    outputs = model(**inputs)

    # Processar os resultados
    target_sizes = torch.tensor([image.size[::-1]])
    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)[0]

    # Desenhar as caixas delimitadoras na imagem
    draw = ImageDraw.Draw(image)
    num_pessoas = 0

    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
        if score > 0.9 and model.config.id2label[label.item()] == "person":
            box = box.tolist()
            draw.rectangle(box, outline="red", width=3)
            num_pessoas += 1

    return image, f"Número de pessoas: {num_pessoas}"

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

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