Update tasks/image.py

tasks/image.py

@@ -116,20 +116,24 @@ async def evaluate_image(request: ImageEvaluationRequest):
         if has_smoke:
             image_true_boxes = parse_boxes(annotation)
             true_boxes_list.append(image_true_boxes)
-
+    
             # Perform YOLO inference
             results = yolo_model.predict(image)
-            if len(results[0].boxes):  # If predictions exist
-                pred_box = results[0].boxes.xywh[0].cpu().numpy().tolist()  # First box in YOLO format
-                pred_boxes.append(pred_box)
+    
+            # 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_boxes.append([])  # No prediction for this image
+                pred_box = []  # No prediction for this image
         else:
             true_boxes_list.append([])
-            pred_boxes.append([])
-
+            pred_box = []  # No ground truth or predictions for this image
+    
+        pred_boxes.append(pred_box)
+    
         # Classification: If predictions exist, assume smoke is present
         predictions.append(1 if len(results[0].boxes) > 0 else 0)
+
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE