[feat] add yolo nano inference baseline

requirements.txt

@@ -1,3 +1,5 @@
+--extra-index-url https://download.pytorch.org/whl/cu124
+
 fastapi>=0.68.0
 uvicorn>=0.15.0
 codecarbon>=2.3.1
@@ -7,4 +9,7 @@ pydantic>=1.10.0
 python-dotenv>=1.0.0
 gradio>=4.0.0
 requests>=2.31.0
-librosa==0.10.2.post1
+librosa==0.10.2.post1
+ultralytics==8.3.59
+torch==2.5.1
+torchvision==0.20.1

tasks/image.py

@@ -14,7 +14,7 @@ load_dotenv()
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "Frugal Object Detector for forest fires"
 ROUTE = "/image"
 
 def parse_boxes(annotation_string):
@@ -73,9 +73,7 @@ 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
+    Current Model: Yolo11 nano
     
     Metrics:
     - Classification accuracy: Whether an image contains smoke or not
@@ -98,8 +96,17 @@ async def evaluate_image(request: ImageEvaluationRequest):
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE CODE HERE
     # Update the code below to replace the random baseline with your model inference
-    #--------------------------------------------------------------------------------------------   
-    
+    #--------------------------------------------------------------------------------------------  
+    from pathlib import Path
+    from ultralytics import YOLO
+    import torch
+
+    # Load model
+    model_path = "models"
+    model_name = "best.pt"
+    model = YOLO(Path(model_path, model_name))
+    threshold = 0.14
+
     predictions = []
     true_labels = []
     pred_boxes = []
@@ -110,26 +117,28 @@ async def evaluate_image(request: ImageEvaluationRequest):
         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 prediction
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        results = model.predict(example["image"], device=device, conf=threshold, verbose=False)[0] # index 0 since we predict on one image at a time
+        if results.boxes.cls.numel()!=0:
+            # This means a fire was detected, hence we append 1
+            pred_has_smoke = 1
+        else:
+            pred_has_smoke = 0
         predictions.append(int(pred_has_smoke))
         
-        # If there's a true box, parse it and make random box prediction
+        # If there's a true box, parse it and add 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)
-            
-            # 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)
+
+            # Append only one bounding box if at least one fire is detected
+            if results.boxes.cls.numel()!=0:
+                pred_boxes.append(results.boxes[0].xywhn.tolist()[0])
+            else:
+                pred_boxes.append([0,0,0,0])
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE

tasks/models/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08ca51a239f739eab4f3653956abcf303f836e8ea3b9a1c225c85f0cc1d086fa
+size 5443539

tasks/models/pruned.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e5e9ef2d0bbe8e8984d6739ccc2d21045844c2be98425b271090de621042ce8
+size 5470665