Update app.py

app.py

@@ -3,7 +3,6 @@ import gradio as gr
 from ultralytics import ASSETS, YOLO
 
 model = YOLO("best.pt")
-
 def predict_image(img, conf_threshold, iou_threshold):
     results = model.predict(
         source=img,
@@ -14,18 +13,25 @@ def predict_image(img, conf_threshold, iou_threshold):
         imgsz=640,
     )
 
-    # Assuming 'results' has a property or method to get detected labels count
-    # If not directly available, you might need to parse the results accordingly
-    fruits_count = sum(1 for _ in results if _.label == "fruit")  # Example, adjust based on actual results structure
+    # Access the predictions and class names
+    preds = results.pred[0]
+    class_ids = preds[:, -1].int().tolist()  # Get class IDs for all detections
+    class_names = [results.names[i] for i in class_ids]
 
-    # Only handling the last result for simplicity, adjust according to your needs
-    for r in results:
-        im_array = r.plot()
-        im = Image.fromarray(im_array[..., ::-1])
-    
-    caption = f"Detected fruits: {fruits_count}"  # Modify this line according to the actual object you're detecting
-    return im, caption
+    # Count the number of fruits detected
+    fruit_count = sum(1 for name in class_names if 'fruit' in name.lower())  # Adjust 'fruit' to the actual class name if needed
 
+    # Convert the result to an image and draw the count on it
+    im_array = results.render()[0]
+    im = Image.fromarray(im_array)
+    draw = ImageDraw.Draw(im)
+    text = f"Fruits detected: {fruit_count}"
+    text_size = draw.textsize(text)
+    draw.rectangle([(0, im.height - text_size[1]), (text_size[0], im.height)], fill="white")
+    draw.text((0, im.height - text_size[1]), text, fill="black")
+
+    return im
+    
 iface = gr.Interface(
     fn=predict_image,
     inputs=[