Final try

tasks/image.py

@@ -114,36 +114,38 @@ async def evaluate_image(request: ImageEvaluationRequest):
     logging.info(f"Loading model {model_name}")
     model = YOLO(Path(model_path, model_name), task="detect")
     device_name = device("cuda" if is_available() else "cpu")
+    model.to(device_name)
     
+    # Preprocessing the annotations before the loop to avoid repeated parsing
+    annotations = [example.get("annotations", "").strip() for example in test_dataset]
+    true_labels = [int(len(ann) > 0) for ann in annotations]
+
+    # Initialize lists
     predictions = []
-    true_labels = []
+    true_boxes_list = []
     pred_boxes = []
-    true_boxes_list = []  # List of lists, each inner list contains boxes for one image
-    
+
     logging.info(f"Inference start on device: {device_name}")
-    for example in test_dataset:
-        # Parse true annotation (YOLO format: class_id x_center y_center width height)
-        annotation = example.get("annotations", "").strip()
-        has_smoke = len(annotation) > 0
-        true_labels.append(int(has_smoke))
+    for i, example in enumerate(test_dataset):
+        has_smoke = true_labels[i]
+        annotation = annotations[i]
 
         # Make prediction
         results = model.predict(example["image"], device=device_name, conf=THRESHOLD, verbose=False, imgsz=IMGSIZE)[0]
         pred_has_smoke = len(results) > 0
         predictions.append(int(pred_has_smoke))
-        
+
         # If there's a true box, parse it and add box prediction
         if has_smoke:
             # Parse all true boxes from the annotation
             image_true_boxes = parse_boxes(annotation)
             true_boxes_list.append(image_true_boxes)
 
-            # Append only one bounding box if at least one fire is detected
-            # Note that multiple boxes could be appended
-            if results.boxes.cls.numel()!=0:
+            # Append bounding box for the prediction
+            if results.boxes.cls.numel() != 0:
                 pred_boxes.append(results.boxes[0].xywhn.tolist()[0])
             else:
-                pred_boxes.append([0,0,0,0])
+                pred_boxes.append([0, 0, 0, 0])
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE