Optimize submission

tasks/image.py

@@ -18,6 +18,7 @@ from pathlib import Path
 from ultralytics import YOLO
 from torch import device
 from torch.cuda import is_available
+from torch import no_grad
 
 router = APIRouter()
 
@@ -115,35 +116,52 @@ async def evaluate_image(request: ImageEvaluationRequest):
     model = YOLO(Path(model_path, model_name), task="detect")
     device_name = device("cuda" if is_available() else "cpu")
     
+    # Preprocess annotations before the loop
+    preprocessed_annotations = [parse_boxes(example.get("annotations", "").strip()) for example in test_dataset]
+
+    batch_size = 16  # Define a batch size
+    batch_images = []
+    batch_annotations = []
+
     predictions = []
     true_labels = []
     pred_boxes = []
     true_boxes_list = []  # List of lists, each inner list contains boxes for one image
-    
+
     logging.info(f"Inference start on device: {device_name}")
-    for example in test_dataset:
-        # 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 prediction
-        results = model.predict(example["image"], device=device_name, conf=THRESHOLD, verbose=False, half=True, imgsz=IMGSIZE)[0]
-        pred_has_smoke = len(results) > 0
-        predictions.append(int(pred_has_smoke))
-        
-        # 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)
-
-            # Append only one bounding box if at least one fire is detected
-            # Note that multiple boxes could be appended
-            if results.boxes.cls.numel()!=0:
-                pred_boxes.append(results.boxes[0].xywhn.tolist()[0])
-            else:
-                pred_boxes.append([0,0,0,0])
+
+    # Use torch.no_grad() to disable gradient tracking during inference
+    with no_grad():
+        for idx, example in enumerate(test_dataset):
+            batch_images.append(example["image"])
+            batch_annotations.append(preprocessed_annotations[idx])
+
+            # When the batch size is met, or it's the last image, perform inference
+            if (len(batch_images) == batch_size or idx == len(test_dataset) - 1):
+                # Make a prediction for the current batch
+                results = model.predict(batch_images, device=device_name, conf=THRESHOLD, verbose=False, half=True, imgsz=IMGSIZE)[0]
+
+                for batch_idx, result in enumerate(results):
+                    annotation = batch_annotations[batch_idx]
+                    has_smoke = len(annotation) > 0
+                    true_labels.append(int(has_smoke))
+
+                    pred_has_smoke = len(result) > 0
+                    predictions.append(int(pred_has_smoke))
+
+                    if has_smoke:
+                        true_boxes_list.append(annotation)
+
+                    # Handle prediction boxes for each image in the batch
+                    if result.boxes.cls.numel() != 0:
+                        pred_boxes.append(result.boxes[0].xywhn.tolist()[0])
+                    else:
+                        pred_boxes.append([0, 0, 0, 0])
+
+                # Clear the batch after processing
+                batch_images.clear()
+                batch_annotations.clear()
+
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE