Update app.py

app.py

@@ -3,6 +3,8 @@ 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,
@@ -13,25 +15,13 @@ def predict_image(img, conf_threshold, iou_threshold):
         imgsz=640,
     )
 
-    # 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]
+    for r in results:
+        im_array = r.plot()
+        im = Image.fromarray(im_array[..., ::-1])
 
-    # 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
+    return im
 
-    # 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=[
@@ -39,10 +29,11 @@ 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"), gr.Textbox(label="Caption")],
+    outputs=gr.Image(type="pil", label="Result"),
     title="My Yield | 🌱",
     description="Estimate the amount of plants per year",
+
 )
 
 if __name__ == '__main__':
-    iface.launch()
+    iface.launch()