Update tasks/image.py

tasks/image.py

@@ -122,7 +122,7 @@ async def evaluate_image(request: ImageEvaluationRequest):
         # Make prediction with model
         image = example['image']
         with torch.inference_mode():
-            pred = model(image)[0]
+            pred = model(image, conf=0.03)[0]
         smoke_detected = pred.boxes.xywhn.shape[0] > 0
         predictions.append(int(smoke_detected))
         
@@ -132,7 +132,8 @@ async def evaluate_image(request: ImageEvaluationRequest):
             # Parse all true boxes from the annotation
             image_true_boxes = parse_boxes(annotation)
             true_boxes_list.append(image_true_boxes)
-            pred_boxes.append(pred.boxes.xywhn.tolist()[0])
+            box_idx = pred.boxes.conf.argmax().item()
+            pred_boxes.append(pred.boxes.xywhn.tolist()[box_idx])
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE