counting added

app.py

@@ -1,12 +1,9 @@
 # app.py
 import gradio as gr
 from PIL import Image, ImageDraw, ImageFont
-import torch
 from ultralytics import YOLO
 import numpy as np
 import os
-from PIL import __version__ as PIL_VERSION
-print(f"Pillow version: {PIL_VERSION}")
 
 MODEL_PATH = "model/231220_detect_lr_0001_640_brightness.pt"
 
@@ -17,7 +14,8 @@ print("YOLO model loaded.")
 
 def detect_plastic_pellets(input_image, threshold=0.5):
     """
-    Perform plastic pellet detection using our customized model.
+    Perform plastic pellet detection using our customized YOLO model.
+    Returns the processed image and the number of detections.
     """
     if input_image is None:
         error_image = Image.new('RGB', (500, 100), color=(255, 0, 0))
@@ -27,7 +25,7 @@ def detect_plastic_pellets(input_image, threshold=0.5):
         except IOError:
             font = ImageFont.load_default()
         draw.text((10, 40), "Please upload a valid image.", fill=(255, 255, 255), font=font)
-        return error_image
+        return error_image, 0  # Returning 0 detections
 
     try:
         print("Starting detection with threshold:", threshold)
@@ -42,6 +40,7 @@ def detect_plastic_pellets(input_image, threshold=0.5):
             font = ImageFont.load_default()
 
         detection_made = False
+        detection_count = 0  # Initialize detection count
 
         for result in results:
             for box in result.boxes:
@@ -58,17 +57,19 @@ def detect_plastic_pellets(input_image, threshold=0.5):
 
                 label = f"{name} {confidence:.2f}"
                 text_width, text_height = font.getbbox(label)[2:]
-                draw.rectangle(((x1, y1 - text_height), (x1 + text_width, y1)), fill=color)
-                draw.text((x1, y1 - text_height), label, fill=(255, 255, 255), font=font)
+                # Ensure text does not go above the image
+                text_y = max(y1 - text_height, 0)
+                draw.rectangle(((x1, text_y), (x1 + text_width, y1)), fill=color)
+                draw.text((x1, text_y), label, fill=(255, 255, 255), font=font)
 
                 detection_made = True
+                detection_count += 1  # Increment detection count
 
         if not detection_made:
             draw.text((10, 10), "No plastic pellets detected.", fill=(255, 0, 0), font=font)
-            return input_image  
-
-        print("Detection completed.")
-        return input_image
+        
+        print("Detection completed. Total detections:", detection_count)
+        return input_image, detection_count
 
     except Exception as e:
         print(f"Detection error: {str(e)}")
@@ -79,7 +80,7 @@ def detect_plastic_pellets(input_image, threshold=0.5):
         except IOError:
             font = ImageFont.load_default()
         draw.text((10, 40), f"Error: {str(e)}", fill=(255, 255, 255), font=font)
-        return error_image
+        return error_image, 0  # Returning 0 detections on error
 
 def main():
     with gr.Blocks(css=".gradio-container {max-width: 800px}") as demo:
@@ -93,10 +94,17 @@ def main():
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(type="pil", label="🌊 Upload or Select Beach Image", interactive=True)
-                examples = ['images/image1.bmp', 'images/image2.bmp', 'images/image3.bmp', 'images/image4.bmp', 'images/image5.bmp', 'images/image6.bmp']
+                examples = [
+                    'images/image1.bmp',
+                    'images/image2.bmp',
+                    'images/image3.bmp',
+                    'images/image4.bmp',
+                    'images/image5.bmp',
+                    'images/image6.bmp'
+                ]
                 gr.Examples(examples=examples, inputs=input_image, label="Or choose one of these images")
                 
-                # Add a slider for confidence threshold
+                # Slider for confidence threshold
                 confidence_threshold = gr.Slider(
                     minimum=0.0,
                     maximum=1.0,
@@ -109,7 +117,17 @@ def main():
                 submit_button = gr.Button("🔍 Detect Plastic Pellets")
 
             with gr.Column():
-                output_image = gr.Image(type="pil", label="✅ Detection Result", interactive=False, show_download_button=True)
+                output_image = gr.Image(
+                    type="pil",
+                    label="✅ Detection Result",
+                    interactive=False,
+                    show_download_button=True
+                )
+                detection_count = gr.Text(
+                    value="Detections: 0",
+                    label="🔢 Number of Detections",
+                    interactive=False
+                )
 
         gr.Markdown(
             """
@@ -121,7 +139,7 @@ def main():
         submit_button.click(
             fn=detect_plastic_pellets,
             inputs=[input_image, confidence_threshold],
-            outputs=output_image,
+            outputs=[output_image, detection_count],  
             api_name="detect",
             show_progress=True  
         )
@@ -129,4 +147,4 @@ def main():
     demo.launch()
 
 if __name__ == "__main__":
-        main()
+    main()