Update tasks/image.py

tasks/image.py

@@ -18,7 +18,7 @@ router = APIRouter()
 DESCRIPTION = "YOLO Smoke Detection"
 ROUTE = "/image"
 
-yolo_model = YOLO("best.pt")
+yolo_model = YOLO("yolo11_tr_20012025_frugalAIchal.pt")
 
 def parse_boxes(annotation_string):
     """Parse multiple boxes from a single annotation string.
@@ -104,57 +104,50 @@ async def evaluate_image(request: ImageEvaluationRequest):
     #--------------------------------------------------------------------------------------------   
     predictions = []
     true_labels = []
-    pred_boxes = []  # Flattened list of predicted boxes
-    true_boxes_list = []  # Flattened list of ground truth boxes
+    pred_boxes = []  
+    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(1 if has_smoke else 0)
     
-        # Parse ground truth boxes if smoke is present
+
         if has_smoke:
             image_true_boxes = parse_boxes(annotation)
             if image_true_boxes:
                 true_boxes_list.append(image_true_boxes)
             else:
-                true_boxes_list.append([])  # Add empty list if parsing fails
+                true_boxes_list.append([])
         else:
-            true_boxes_list.append([])  # Add empty list for no ground truth smoke
+            true_boxes_list.append([])
     
-        # Perform YOLO inference
-        results = yolo_model .predict(image, verbose=False)
+        results = yolo_model .predict(image, verbose=False)    # INFERENCE - prediction
         
-        # Extract predicted box if predictions exist
         if len(results[0].boxes):
             pred_box = results[0].boxes.xywhn[0].cpu().numpy().tolist()
-            predictions.append(1)  # Predicted smoke
+            predictions.append(1)  
             pred_boxes.append(pred_box)
         else:
-            predictions.append(0)  # No smoke predicted
-            pred_boxes.append([])  # Add empty list if no prediction
+            predictions.append(0)  
+            pred_boxes.append([])  
 
-    # 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
+    for true_boxes, pred_boxes_entry in zip(true_boxes_list, pred_boxes):  # Only see when annotation(s) is/are both on true label and prediction
+        if true_boxes and pred_boxes_entry:  
             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
     
-       
-
-    
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
     #--------------------------------------------------------------------------------------------