Update app.py

app.py

@@ -1,34 +1,43 @@
 import gradio as gr
-from PIL import Image, ImageDraw
+from transformers import DetrImageProcessor, DetrForObjectDetection
 import torch
-from facenet_pytorch import MTCNN
+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)
 
-# Carregar o modelo de detecção de rostos
-mtcnn = MTCNN(keep_all=True)
+    # Processar os resultados
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)[0]
 
-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
+    num_pessoas = 0
 
-    if boxes is not None:
-        for box in boxes:
-            draw.rectangle(box.tolist(), outline="red", width=3)
-            num_faces += 1
+    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 rostos: {num_faces}"
+    return image, f"Número de pessoas: {num_pessoas}"
 
 # Criar a interface do Gradio
 iface = gr.Interface(
-    fn=detect_faces, 
+    fn=detect_people, 
     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."
+    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()
+iface.launch()