Update tasks/image.py

tasks/image.py

@@ -19,6 +19,13 @@ router = APIRouter()
 DESCRIPTION = "YOLOv11"
 ROUTE = "/image"
 
+def collate_fn(batch):
+    """Prepare a batch of examples."""
+    images = [example["image"] for example in batch]
+    annotations = [example["annotations"].strip() for example in batch]
+    return images, annotations
+
+
 def parse_boxes(annotation_string):
     """Parse multiple boxes from a single annotation string.
     Each box has 5 values: class_id, x_center, y_center, width, height"""
@@ -83,100 +90,75 @@ def get_boxes_list(predictions):
              description=DESCRIPTION)
 async def evaluate_image(request: ImageEvaluationRequest):
     """
-    Evaluate image classification and object detection for forest fire smoke.
-    
-    Current Model: Random Baseline
-    - Makes random predictions for both classification and bounding boxes
-    - Used as a baseline for comparison
-    
-    Metrics:
-    - Classification accuracy: Whether an image contains smoke or not
-    - Object Detection accuracy: IoU (Intersection over Union) for smoke bounding boxes
+    Evaluate image classification and object detection for forest fire smoke using batched inference.
     """
-    # Get space info
-    username, space_url = get_space_info()
-    
     # Load and prepare the dataset
     dataset = load_dataset(request.dataset_name, token=os.getenv("HF_TOKEN"))
-    
-    # Split dataset
     train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
     test_dataset = train_test["test"]
-    
-    # Start tracking emissions
+
+    # Load YOLO model
+    model_path = "best.pt"
+    model = YOLO(model_path)
+    model.eval()
+
+    # Set up DataLoader for batched processing
+    batch_size = 8
+    dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
+
+    # Initialize variables for evaluation
     tracker.start()
     tracker.start_task("inference")
-    
-    #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE CODE HERE
-    # Update the code below to replace the random baseline with your model inference
-    #--------------------------------------------------------------------------------------------   
-    
-    
-    PATH_TO_MODEL = f"best.pt"
-    model = load_model(PATH_TO_MODEL)
-    
-    print(f"Model info: {model.info()}")
-    predictions = []
+
     true_labels = []
+    predictions = []
     pred_boxes = []
-    true_boxes_list = []  # List of lists, each inner list contains boxes for one image
-    
-    n_examples = len(test_dataset)
-    for i, example in enumerate(test_dataset):
-        print(f"Running {i+1} of {n_examples}")
-        # Parse true annotation (YOLO format: class_id x_center y_center width height)
-        annotation = example.get("annotations", "").strip()
-        has_smoke = len(annotation) > 0
-        true_labels.append(int(has_smoke))
-        
-        model_preds = model(example['image'])[0]
-        pred_has_smoke = len(model_preds) > 0
-        predictions.append(int(pred_has_smoke))
+    true_boxes_list = []
+
+    for batch_idx, (images, annotations) in enumerate(dataloader):
+        print(f"Processing batch {batch_idx + 1}")
+
+        # Parse true labels and boxes
+        batch_true_labels = []
+        batch_true_boxes_list = []
+        for annotation in annotations:
+            has_smoke = len(annotation) > 0
+            batch_true_labels.append(int(has_smoke))
+            true_boxes = parse_boxes(annotation) if has_smoke else []
+            batch_true_boxes_list.append(true_boxes)
         
-        # If there's a true box, parse it and make random box prediction
-        if has_smoke:
-            
-            # Parse all true boxes from the annotation
-            image_true_boxes = parse_boxes(annotation)
-            true_boxes_list.append(image_true_boxes)
-            
-            try:
-                pred_box_list = get_boxes_list(model_preds)[0] # With one bbox to start with (as in the random baseline)
-            except:
-                print("No boxes found")
-                pred_box_list = [0, 0, 0, 0] # Hacky way to make sure that compute_max_iou doesn't fail
-            pred_boxes.append(pred_box_list)
+        true_labels.extend(batch_true_labels)
+        true_boxes_list.extend(batch_true_boxes_list)
 
+        # YOLO batch inference
+        batch_predictions = model(images)
+
+        # Parse predictions for smoke detection and bounding boxes
+        batch_predictions_classes = [1 if len(pred.boxes) > 0 else 0 for pred in batch_predictions]
+        batch_pred_boxes = [get_boxes_list(pred)[0] if len(pred.boxes) > 0 else [0, 0, 0, 0] for pred in batch_predictions]
+
+        predictions.extend(batch_predictions_classes)
+        pred_boxes.extend(batch_pred_boxes)
 
-    
-    #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE STOPS HERE
-    #--------------------------------------------------------------------------------------------   
-    
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
-    
+
     # Calculate classification metrics
     classification_accuracy = accuracy_score(true_labels, predictions)
     classification_precision = precision_score(true_labels, predictions)
     classification_recall = recall_score(true_labels, predictions)
-    
+
     # Calculate mean IoU for object detection (only for images with smoke)
-    # For each image, we compute the max IoU between the predicted box and all true boxes
-    ious = []
-    for true_boxes, pred_box in zip(true_boxes_list, pred_boxes):
-        max_iou = compute_max_iou(true_boxes, pred_box)
-        ious.append(max_iou)
-    
+    ious = [compute_max_iou(true_boxes, pred_box) for true_boxes, pred_box in zip(true_boxes_list, pred_boxes)]
     mean_iou = float(np.mean(ious)) if ious else 0.0
-    
+
     # Prepare results dictionary
+    username, space_url = get_space_info()
     results = {
         "username": username,
         "space_url": space_url,
         "submission_timestamp": datetime.now().isoformat(),
-        "model_description": DESCRIPTION,
+        "model_description": "YOLOv11",
         "classification_accuracy": float(classification_accuracy),
         "classification_precision": float(classification_precision),
         "classification_recall": float(classification_recall),
@@ -184,12 +166,125 @@ async def evaluate_image(request: ImageEvaluationRequest):
         "energy_consumed_wh": emissions_data.energy_consumed * 1000,
         "emissions_gco2eq": emissions_data.emissions * 1000,
         "emissions_data": clean_emissions_data(emissions_data),
-        "api_route": ROUTE,
+        "api_route": "/image",
         "dataset_config": {
             "dataset_name": request.dataset_name,
             "test_size": request.test_size,
             "test_seed": request.test_seed
         }
     }
+
+    return results
+
+# async def evaluate_image(request: ImageEvaluationRequest):
+#     """
+#     Evaluate image classification and object detection for forest fire smoke.
+    
+#     Current Model: Random Baseline
+#     - Makes random predictions for both classification and bounding boxes
+#     - Used as a baseline for comparison
+    
+#     Metrics:
+#     - Classification accuracy: Whether an image contains smoke or not
+#     - Object Detection accuracy: IoU (Intersection over Union) for smoke bounding boxes
+#     """
+#     # Get space info
+#     username, space_url = get_space_info()
+    
+#     # Load and prepare the dataset
+#     dataset = load_dataset(request.dataset_name, token=os.getenv("HF_TOKEN"))
+    
+#     # Split dataset
+#     train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
+#     test_dataset = train_test["test"]
+    
+#     # Start tracking emissions
+#     tracker.start()
+#     tracker.start_task("inference")
+    
+#     #--------------------------------------------------------------------------------------------
+#     # YOUR MODEL INFERENCE CODE HERE
+#     # Update the code below to replace the random baseline with your model inference
+#     #--------------------------------------------------------------------------------------------   
+    
+    
+#     PATH_TO_MODEL = f"best.pt"
+#     model = load_model(PATH_TO_MODEL)
+    
+#     print(f"Model info: {model.info()}")
+#     predictions = []
+#     true_labels = []
+#     pred_boxes = []
+#     true_boxes_list = []  # List of lists, each inner list contains boxes for one image
+    
+#     n_examples = len(test_dataset)
+#     for i, example in enumerate(test_dataset):
+#         print(f"Running {i+1} of {n_examples}")
+#         # Parse true annotation (YOLO format: class_id x_center y_center width height)
+#         annotation = example.get("annotations", "").strip()
+#         has_smoke = len(annotation) > 0
+#         true_labels.append(int(has_smoke))
+        
+#         model_preds = model(example['image'])[0]
+#         pred_has_smoke = len(model_preds) > 0
+#         predictions.append(int(pred_has_smoke))
+        
+#         # If there's a true box, parse it and make random box prediction
+#         if has_smoke:
+            
+#             # Parse all true boxes from the annotation
+#             image_true_boxes = parse_boxes(annotation)
+#             true_boxes_list.append(image_true_boxes)
+            
+#             try:
+#                 pred_box_list = get_boxes_list(model_preds)[0] # With one bbox to start with (as in the random baseline)
+#             except:
+#                 print("No boxes found")
+#                 pred_box_list = [0, 0, 0, 0] # Hacky way to make sure that compute_max_iou doesn't fail
+#             pred_boxes.append(pred_box_list)
+
+
+    
+#     #--------------------------------------------------------------------------------------------
+#     # YOUR MODEL INFERENCE STOPS HERE
+#     #--------------------------------------------------------------------------------------------   
+    
+#     # Stop tracking emissions
+#     emissions_data = tracker.stop_task()
+    
+#     # Calculate classification metrics
+#     classification_accuracy = accuracy_score(true_labels, predictions)
+#     classification_precision = precision_score(true_labels, predictions)
+#     classification_recall = recall_score(true_labels, predictions)
+    
+#     # Calculate mean IoU for object detection (only for images with smoke)
+#     # For each image, we compute the max IoU between the predicted box and all true boxes
+#     ious = []
+#     for true_boxes, pred_box in zip(true_boxes_list, pred_boxes):
+#         max_iou = compute_max_iou(true_boxes, pred_box)
+#         ious.append(max_iou)
+    
+#     mean_iou = float(np.mean(ious)) if ious else 0.0
+    
+#     # Prepare results dictionary
+#     results = {
+#         "username": username,
+#         "space_url": space_url,
+#         "submission_timestamp": datetime.now().isoformat(),
+#         "model_description": DESCRIPTION,
+#         "classification_accuracy": float(classification_accuracy),
+#         "classification_precision": float(classification_precision),
+#         "classification_recall": float(classification_recall),
+#         "mean_iou": mean_iou,
+#         "energy_consumed_wh": emissions_data.energy_consumed * 1000,
+#         "emissions_gco2eq": emissions_data.emissions * 1000,
+#         "emissions_data": clean_emissions_data(emissions_data),
+#         "api_route": ROUTE,
+#         "dataset_config": {
+#             "dataset_name": request.dataset_name,
+#             "test_size": request.test_size,
+#             "test_seed": request.test_seed
+#         }
+#     }
     
-    return results 
+#     return results