Update app.py

app.py

@@ -29,7 +29,7 @@ model = YOLO(model_path).to(device)
 
 # Define the detection function for Gradio
 @spaces.GPU  # Decorator to allocate GPU for ZeroGPU-enabled Spaces
-def detect_objects(image: np.ndarray) -> Image.Image:
+def detect_objects(image: np.ndarray) -> (Image.Image, str):
     # Ensure the image is in BGR format if provided by PIL (Gradio gives us an RGB image)
     if image.shape[-1] == 3:
         image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
@@ -57,6 +57,9 @@ def detect_objects(image: np.ndarray) -> Image.Image:
     # Apply Non-Maximum Suppression (NMS) to the detections to avoid duplicate boxes
     detections = detections.with_nms(threshold=NMS_THRESHOLD, class_agnostic=False)
 
+    # Count total detections after NMS
+    total_detections = len(detections)
+
     # Initialize an annotator for bounding boxes with specified color and thickness
     box_annotator = sv.OrientedBoxAnnotator(color=ANNOTATION_COLOR, thickness=ANNOTATION_THICKNESS)
 
@@ -65,11 +68,13 @@ def detect_objects(image: np.ndarray) -> Image.Image:
 
     # Convert annotated image to RGB for Gradio display (PIL expects RGB)
     annotated_img_rgb = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
-    return Image.fromarray(annotated_img_rgb)
+
+    # Return the annotated image and the total count of detections
+    return Image.fromarray(annotated_img_rgb), f"Total Detections: {total_detections}"
 
 # Reset function for Gradio UI
 def gradio_reset():
-    return gr.update(value=None), gr.update(value=None)
+    return gr.update(value=None), gr.update(value=None), gr.update(value="")
 
 # Set up Gradio interface
 with gr.Blocks() as demo:
@@ -83,6 +88,7 @@ with gr.Blocks() as demo:
             
         with gr.Column():
             output_img = gr.Image(label="Detection Result", interactive=False)
+            detection_count = gr.Textbox(label="Detection Summary", interactive=False)
     
     # Add Examples section with images from the root directory
     with gr.Accordion("Select an Example Image"):
@@ -94,8 +100,8 @@ with gr.Blocks() as demo:
         )
     
     # Define button actions
-    clear.click(gradio_reset, inputs=None, outputs=[input_img, output_img])
-    predict.click(detect_objects, inputs=[input_img], outputs=[output_img])
+    clear.click(gradio_reset, inputs=None, outputs=[input_img, output_img, detection_count])
+    predict.click(detect_objects, inputs=[input_img], outputs=[output_img, detection_count])
 
 # Launch Gradio app
 demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)