Return person count

app.py

@@ -1,13 +1,11 @@
 import gradio as gr
 import PIL.Image as Image
-
 from ultralytics import ASSETS, YOLO
 
 model = YOLO("yolo12n.pt")
 
-
 def predict_image(img, conf_threshold, iou_threshold):
-    """Predicts objects in an image."""
+    """Predicts persons in an image and returns the image with detections and count."""
     results = model.predict(
         source=img,
         conf=conf_threshold,
@@ -21,9 +19,10 @@ def predict_image(img, conf_threshold, iou_threshold):
     for r in results:
         im_array = r.plot()
         im = Image.fromarray(im_array[..., ::-1])
+    
+    person_count = len(results[0].boxes) if results[0].boxes is not None else 0
 
-    return im
-
+    return im, f"Number of persons detected: {person_count}"
 
 iface = gr.Interface(
     fn=predict_image,
@@ -32,13 +31,12 @@ iface = gr.Interface(
         gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
         gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
     ],
-    outputs=gr.Image(type="pil", label="Result"),
-    title="Image",
-    description="Upload images for inference",
-    # examples=[
-    #     [ASSETS / "bus.jpg", 0.25, 0.45],
-    #     [ASSETS / "zidane.jpg", 0.25, 0.45],
-    # ],
+    outputs=[
+        gr.Image(type="pil", label="Result"),
+        gr.Textbox(label="Person Count")
+    ],
+    title="Image Person Detection",
+    description="Upload images to detect persons and get a count",
 )
 
 if __name__ == "__main__":