Update tasks/image.py

tasks/image.py

@@ -108,31 +108,37 @@ async def evaluate_image(request: ImageEvaluationRequest):
     true_boxes_list = []
 
     for example in test_dataset:
+        # 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))
-        
+    
+        # Parse ground truth boxes if smoke is present
         if has_smoke:
             image_true_boxes = parse_boxes(annotation)
             true_boxes_list.append(image_true_boxes)
+        else:
+            true_boxes_list.append([])  # No ground truth boxes
     
-            # Perform YOLO inference
-            results = yolo_model.predict(image)
+        # Perform YOLO inference
+        results = yolo_model.predict(image)
     
-            # Extract predicted box if it exists
-            if len(results[0].boxes):  # Check if predictions exist
-                pred_box = results[0].boxes.xywh[0].cpu().numpy().tolist()  # Take the first box
-            else:
-                pred_box = []  # No prediction for this 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
+            predictions.append(1)  # Predicted smoke
         else:
-            true_boxes_list.append([])
-            pred_box = []  # No ground truth or predictions for this image
+            pred_box = []  # No prediction for this image
+            predictions.append(0)  # Predicted no smoke
     
+        # Append the predicted box (empty if no predictions)
         pred_boxes.append(pred_box)
+
     
-        # Classification: If predictions exist, assume smoke is present
-        predictions.append(1 if len(results[0].boxes) > 0 else 0)
+       
 
     
     #--------------------------------------------------------------------------------------------