Update tasks/image.py

tasks/image.py

@@ -104,9 +104,10 @@ async def evaluate_image(request: ImageEvaluationRequest):
     #--------------------------------------------------------------------------------------------   
     predictions = []
     true_labels = []
-    pred_boxes = []
-    true_boxes_list = []
-
+    pred_boxes = []  # Flattened list of predicted boxes
+    true_boxes_list = []  # Flattened list of ground truth boxes
+    
+    
     for example in test_dataset:
         # Extract image and annotations
         image = example["image"]
@@ -119,19 +120,19 @@ async def evaluate_image(request: ImageEvaluationRequest):
         # Parse ground truth boxes if smoke is present
         if has_smoke:
             image_true_boxes = parse_boxes(annotation)
-            true_boxes_list.append(image_true_boxes)
+            true_boxes_list.append(image_true_boxes if image_true_boxes else [])
         else:
             true_boxes_list.append([])  # No ground truth boxes
     
         # Perform YOLO inference
-        results = yolo_model.predict(image)
+        results = model.predict(image)
     
         # Extract predicted box if predictions exist
         if len(results[0].boxes):
-            pred_box = results[0].boxes.xywh[0].cpu().numpy().tolist()  # Take the first box
+            pred_box = results[0].boxes.xywh[0].cpu().numpy().tolist()
             predictions.append(1)  # Predicted smoke
         else:
-            pred_box = []  # No prediction for this image
+            pred_box = []  # No predictions for this image
             predictions.append(0)  # Predicted no smoke
     
         # Append the predicted box (empty if no predictions)