Update tasks/image.py

tasks/image.py

@@ -112,32 +112,48 @@ async def evaluate_image(request: ImageEvaluationRequest):
         # Extract image and annotations
         image = example["image"]
         annotation = example.get("annotations", "").strip()
-    
+        
         # Determine if ground truth smoke is present
         has_smoke = len(annotation) > 0
-        true_labels.append(int(has_smoke))
+        true_labels.append(1 if has_smoke else 0)
     
         # Parse ground truth boxes if smoke is present
         if has_smoke:
             image_true_boxes = parse_boxes(annotation)
-            true_boxes_list.append(image_true_boxes if image_true_boxes else [])
+            if image_true_boxes:
+                true_boxes_list.append(image_true_boxes)
+            else:
+                true_boxes_list.append([])  # Add empty list if parsing fails
         else:
-            true_boxes_list.append([])  # No ground truth boxes
+            true_boxes_list.append([])  # Add empty list for no ground truth smoke
     
         # Perform YOLO inference
         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()
+            pred_box = results[0].boxes.xywhn[0].cpu().numpy().tolist()
             predictions.append(1)  # Predicted smoke
+            pred_boxes.append(pred_box)
         else:
-            pred_box = []  # No predictions for this image
-            predictions.append(0)  # Predicted no smoke
+            predictions.append(0)  # No smoke predicted
+            pred_boxes.append([])  # Add empty list if no prediction
     
-        # Append the predicted box (empty if no predictions)
-        pred_boxes.append(pred_box)
-
+    
+    # Filter out entries with empty boxes
+    filtered_true_boxes_list = []
+    filtered_pred_boxes = []
+    
+    
+    for true_boxes, pred_boxes_entry in zip(true_boxes_list, pred_boxes):
+        if true_boxes and pred_boxes_entry:  # Keep only if neither is empty
+            filtered_true_boxes_list.append(true_boxes)
+            filtered_pred_boxes.append(pred_boxes_entry)
+    
+    
+    # Replace the original lists with the filtered ones
+    true_boxes_list = filtered_true_boxes_list
+    pred_boxes = filtered_pred_boxes