Update tasks/image.py

tasks/image.py

@@ -5,6 +5,7 @@ import numpy as np
 from sklearn.metrics import accuracy_score, precision_score, recall_score
 import random
 import os
+from ultralytics import YOLO
 
 from .utils.evaluation import ImageEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
@@ -14,7 +15,7 @@ load_dotenv()
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = f"MountAIn Small model 640"
 ROUTE = "/image"
 
 def parse_boxes(annotation_string):
@@ -67,6 +68,17 @@ def compute_max_iou(true_boxes, pred_box):
         max_iou = max(max_iou, iou)
     return max_iou
 
+def load_model(path_to_model, model_type="YOLO"):
+    if model_type == "YOLO":
+        model = YOLO(path_to_model)
+    else:
+        raise NotImplementedError
+    return model
+
+def get_boxes_list(predictions):
+    return [box.tolist() for box in predictions.boxes.xywhn]
+
+
 @router.post(ROUTE, tags=["Image Task"],
              description=DESCRIPTION)
 async def evaluate_image(request: ImageEvaluationRequest):
@@ -99,20 +111,28 @@ async def evaluate_image(request: ImageEvaluationRequest):
     # YOUR MODEL INFERENCE CODE HERE
     # Update the code below to replace the random baseline with your model inference
     #--------------------------------------------------------------------------------------------   
+
+    PATH_TO_MODEL = f"models/best-mountain-s.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
     
-    for example in test_dataset:
+    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))
         
-        # Make random classification prediction
-        pred_has_smoke = random.random() > 0.5
+        # Make model prediction
+        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
@@ -121,15 +141,12 @@ async def evaluate_image(request: ImageEvaluationRequest):
             image_true_boxes = parse_boxes(annotation)
             true_boxes_list.append(image_true_boxes)
             
-            # For baseline, make one random box prediction per image
-            # In a real model, you might want to predict multiple boxes
-            random_box = [
-                random.random(),  # x_center
-                random.random(),  # y_center
-                random.random() * 0.5,  # width (max 0.5)
-                random.random() * 0.5   # height (max 0.5)
-            ]
-            pred_boxes.append(random_box)
+            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