Optimize submission

tasks/image.py

@@ -106,6 +106,7 @@ async def evaluate_image(request: ImageEvaluationRequest):
     #--------------------------------------------------------------------------------------------  
 
     THRESHOLD = 0.18
+    IMGSIZE = 1280
 
     # Load model
     model_path = Path("tasks", "models")
@@ -113,8 +114,7 @@ async def evaluate_image(request: ImageEvaluationRequest):
     logging.info(f"Loading model {model_name}")
     model = YOLO(Path(model_path, model_name), task="detect")
     device_name = device("cuda" if is_available() else "cpu")
-    IMGSIZE = 1280
-
+    
     predictions = []
     true_labels = []
     pred_boxes = []
@@ -128,7 +128,7 @@ async def evaluate_image(request: ImageEvaluationRequest):
         true_labels.append(int(has_smoke))
 
         # Make prediction
-        results = model.predict(example["image"], device=device_name, conf=THRESHOLD, verbose=True, imgsz=IMGSIZE)[0]
+        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))
         
@@ -139,6 +139,7 @@ async def evaluate_image(request: ImageEvaluationRequest):
             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: