Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from ultralytics import YOLO
-import numpy as np
 from PIL import Image
+from collections import defaultdict
 
 # load model
 model = YOLO("best.pt")
@@ -10,31 +10,40 @@ def predict(image):
     results = model(image)
     r = results[0]
     
-    # Get detection results
-    detections = []
+    # Group detections by class
+    class_confidences = defaultdict(list)
+    
     if len(r.boxes) > 0:
         for box in r.boxes:
             class_id = int(box.cls[0])
             class_name = model.names[class_id]
-            confidence = float(box.conf[0])
-            detections.append({
-                "class": class_name,
-                "confidence": round(confidence * 100, 2)
-            })
+            confidence = float(box.conf[0]) * 100
+            class_confidences[class_name].append(confidence)
     
-    # Sort by confidence (highest first)
-    if detections:
-        detections.sort(key=lambda x: x["confidence"], reverse=True)
+    if class_confidences:
+        lines = []
+        # Sort classes by their max confidence (highest first)
+        sorted_classes = sorted(
+            class_confidences.items(),
+            key=lambda x: max(x[1]),
+            reverse=True
+        )
+        
+        for class_name, confidences in sorted_classes:
+            count = len(confidences)
+            avg_conf = round(sum(confidences) / count, 2)
+            max_conf = round(max(confidences), 2)
+            
+            if count > 1:
+                lines.append(f"• {class_name} ({count}x) — avg: {avg_conf}% | best: {max_conf}%")
+            else:
+                lines.append(f"• {class_name} — {max_conf}%")
         
-        # Build result text with ALL detections
-        lines = [f"Detected {len(detections)} object(s):\n"]
-        for i, d in enumerate(detections, 1):
-            lines.append(f"{i}. {d['class']}: {d['confidence']}%")
-        result_text = "\n".join(lines)
+        total = sum(len(v) for v in class_confidences.values())
+        result_text = f"Detected {total} object(s) in {len(class_confidences)} class(es):\n\n" + "\n".join(lines)
     else:
         result_text = "No detection"
     
-    # Return annotated image AND the detection text
     output_image = r.plot()
     return Image.fromarray(output_image), result_text
 
@@ -43,7 +52,7 @@ demo = gr.Interface(
     inputs=gr.Image(type="pil"),
     outputs=[
         gr.Image(type="pil", label="Detection"),
-        gr.Textbox(label="Result", lines=10)  # taller box to show all results
+        gr.Textbox(label="Result", lines=10)
     ],
     title="YOLOv8 Detection",
     description="Upload an image and detect objects"