Update tasks/image.py

tasks/image.py

@@ -5,6 +5,7 @@ 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,7 +109,7 @@ async def evaluate_image(request: ImageEvaluationRequest):
     true_boxes_list = []  # Flattened list of ground truth boxes
     
     
-    for example in test_dataset:
+    for example in tqdm(test_dataset, desc="Processing test dataset")
         # Extract image and annotations
         image = example["image"]
         annotation = example.get("annotations", "").strip()
@@ -128,7 +129,7 @@ async def evaluate_image(request: ImageEvaluationRequest):
             true_boxes_list.append([])  # Add empty list for no ground truth smoke
     
         # Perform YOLO inference
-        results = yolo_model .predict(image)
+        results = yolo_model .predict(image, verbose=False)
         
         # Extract predicted box if predictions exist
         if len(results[0].boxes):