Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 import torchvision
 from torchvision import transforms
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageDraw
 
 # Use GPU if available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -33,21 +33,37 @@ def count_persons(image):
         with autocast():
             outputs = model(img_tensor)[0]
 
-    # Count persons (label 1 in COCO)
-    person_count = sum(
-        1 for label, score in zip(outputs['labels'], outputs['scores'])
-        if label.item() == 1 and score.item() > 0.65
-    )
+    # Filter for persons (label 1 in COCO) with score > threshold
+    threshold = 0.65
+    boxes = outputs['boxes']
+    labels = outputs['labels']
+    scores = outputs['scores']
+    keep_indices = [
+        i for i, (label, score) in enumerate(zip(labels, scores))
+        if label.item() == 1 and score.item() > threshold
+    ]
 
-    return f"Number of persons detected: {person_count}"
+    person_count = len(keep_indices)
 
-# Gradio interface for image upload
+    # Draw bounding boxes on the image
+    annotated_image = image.convert("RGB")
+    draw = ImageDraw.Draw(annotated_image)
+    for i in keep_indices:
+        box = boxes[i].cpu().numpy()
+        draw.rectangle([ (box[0], box[1]), (box[2], box[3]) ], outline="red", width=2)
+
+    return annotated_image, f"Number of persons detected: {person_count}"
+
+# Gradio interface for image upload and outputs
 demo = gr.Interface(
     fn=count_persons,
     inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Text(label="Person Count"),
+    outputs=[
+        gr.Image(type="pil", label="Annotated Image"),
+        gr.Text(label="Person Count")
+    ],
     title="Person Counter in Image (Fast)",
-    description="Upload an image to count the number of people using a fast MobileNet-based detector. GPU supported."
+    description="Upload an image to count the number of people and see bounding boxes using a fast MobileNet-based detector. GPU supported."
 )
 
 if __name__ == "__main__":