Final submission

tasks/image.py

@@ -116,36 +116,35 @@ async def evaluate_image(request: ImageEvaluationRequest):
     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_boxes_list = []
+    true_labels = []
     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 i, example in enumerate(test_dataset):
-        has_smoke = true_labels[i]
-        annotation = annotations[i]
+    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))
 
         # Make prediction
-        results = model.predict(example["image"], device=device_name, conf=THRESHOLD, verbose=False, imgsz=IMGSIZE)[0]
+        results = model.predict(example["image"], 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 bounding box for the prediction
-            if results.boxes.cls.numel() != 0:
+            # 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:
                 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