Update main.py

main.py

@@ -48,27 +48,25 @@ async def predict(request: ImageRequest):
     json_result = {}
     class_counters = {}
     
-    for box in result.boxes:
+        for box in result.boxes:
         class_id = int(box.cls[0])
         class_name = result.names[class_id]
         x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
 
-        '''
-        Create unique key for each detection
-        EX: if multiple "missing_hole", keys will be:
-        "missing_hole", "missing_hole2", ...
-        '''
+        # Counting occurrences of each defect
         if class_name in class_counters:
             class_counters[class_name] += 1
-            key = f"{class_name}{class_counters[class_name]}"
         else:
             class_counters[class_name] = 1
-            key = class_name
         
+        key = f"{class_name}{class_counters[class_name]}" if class_counters[class_name] > 1 else class_name
         json_result[key] = [x1, y1, x2, y2]
     
-    # Creating statistics summary
-    statistics_summary = {f"{count} {name}{'s' if count > 1 else ''}": count for name, count in class_counters.items()}
+    # Using model's summary output if available
+    if hasattr(result, "summary") and isinstance(result.summary, dict):
+        statistics_summary = result.summary
+    else:
+        statistics_summary = {name: count for name, count in class_counters.items()}
     
     return {"statistics": statistics_summary, "detections": json_result}