Update app.py

app.py

@@ -6,7 +6,7 @@ import numpy as np
 # Load the YOLOv5 model
 model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
 
-# Function to run inference on an image
+# Function to run inference on an image and count objects
 def run_inference(image):
     # Convert the image from PIL format to a format compatible with OpenCV
     image = np.array(image)
@@ -14,19 +14,28 @@ def run_inference(image):
     # Run YOLOv5 inference
     results = model(image)
 
+    # Extract detection results
+    detections = results.pandas().xyxy[0]
+
+    # Count objects by category
+    object_counts = detections['name'].value_counts().to_dict()
+
     # Convert the annotated image from BGR to RGB for display
     annotated_image = results.render()[0]
     annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
 
-    return annotated_image
+    return annotated_image, object_counts
 
 # Create the Gradio interface
 interface = gr.Interface(
     fn=run_inference,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="YOLOv5 Object Detection",
-    description="Upload an image to run YOLOv5 object detection and see the results."
+    outputs=[
+        gr.Image(type="pil"),
+        gr.JSON(label="Object Counts")  # Add JSON output for object counts
+    ],
+    title="YOLOv5 Object Detection with Counts",
+    description="Upload an image to run YOLOv5 object detection, see the annotated results, and view the count of detected objects by category."
 )
 
 # Launch the app