Update app.py

app.py

@@ -2,42 +2,43 @@ from transformers import DetrImageProcessor, DetrForObjectDetection
 from PIL import Image, ImageDraw
 import torch
 import gradio as gr
-import requests
-from io import BytesIO
 
-# Load pre-trained DETR model
+# Load model and processor
 processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
 
-# COCO class index for "person" = 1 (used as proxy for face detection)
-FACE_CLASS_INDEX = 1
+FACE_CLASS_INDEX = 1  # COCO class ID for 'person'
 
 def detect_faces(img: Image.Image):
-    # Prepare input for the model
+    # Make a copy to draw on
+    img_draw = img.copy()
+    draw = ImageDraw.Draw(img_draw)
+
+    # Preprocess and predict
     inputs = processor(images=img, return_tensors="pt")
     outputs = model(**inputs)
 
-    # Get outputs
+    # Get results
     target_sizes = torch.tensor([img.size[::-1]])
-    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.8)[0]
 
-    # Draw bounding boxes
-    draw = ImageDraw.Draw(img)
+    count = 0
     for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-        if label.item() == FACE_CLASS_INDEX:  # 'person'
+        if label.item() == FACE_CLASS_INDEX:
+            count += 1
             box = [round(i, 2) for i in box.tolist()]
-            draw.rectangle(box, outline="green", width=3)
-            draw.text((box[0], box[1]), f"{score:.2f}", fill="green")
+            draw.rectangle(box, outline="lime", width=3)
+            draw.text((box[0], box[1] - 10), f"{score:.2f}", fill="lime")
 
-    return img
+    return img_draw, f"Total Persons Detected: {count}"
 
-# Gradio interface
+# Gradio Interface
 iface = gr.Interface(
     fn=detect_faces,
     inputs=gr.Image(type="pil"),
-    outputs="image",
-    title="Face Detection App (Hugging Face + Gradio)",
-    description="Upload an image and detect faces using facebook/detr-resnet-50 model."
+    outputs=[gr.Image(type="pil"), gr.Text()],
+    title="Person Detection with DETR",
+    description="Uses DETR model to detect people (class 1 - COCO dataset). Note: not specialized for face detection."
 )
 
 iface.launch()