Update tasks/image.py

tasks/image.py

@@ -5,7 +5,6 @@ import numpy as np
 from sklearn.metrics import accuracy_score
 import random
 import os
-from tqdm import tqdm
 
 from ultralytics import YOLO  # Import YOLO
 from .utils.evaluation import ImageEvaluationRequest
@@ -108,8 +107,9 @@ async def evaluate_image(request: ImageEvaluationRequest):
     pred_boxes = []  # Flattened list of predicted boxes
     true_boxes_list = []  # Flattened list of ground truth boxes
     
-    
-    for example in tqdm(test_dataset, desc="Processing test dataset"):
+    i = 0
+    for example in test_dataset:
+        i += 1
         # Extract image and annotations
         image = example["image"]
         annotation = example.get("annotations", "").strip()
@@ -140,7 +140,10 @@ async def evaluate_image(request: ImageEvaluationRequest):
             predictions.append(0)  # No smoke predicted
             pred_boxes.append([])  # Add empty list if no prediction
     
-    
+        if i == 10:
+            break
+
+            
     # Filter out entries with empty boxes
     filtered_true_boxes_list = []
     filtered_pred_boxes = []