Preparing submission

tasks/image.py

@@ -107,27 +107,18 @@ async def evaluate_image(request: ImageEvaluationRequest):
 
     # Load model
     model_path = Path("tasks", "models")
-    # If CUDA is available, load FP16 pytorch
-    # if is_available():
-    # print("CUDA available, loading FP16 pytorch model")
     model_name = "pruned_fp16.pt"
+    print(f"Loading model {model_name}")
     model = YOLO(Path(model_path, model_name), task="detect")
-    # device_name = device("cuda")
     device_name = device("cuda" if is_available() else "cpu")
     IMGSIZE = 1280
-    # # If not, load FP16 ONNX model
-    # else:
-    #     print("CUDA not, available, loading ONNX model")
-    #     model_name = "640_fp16_cpu.onnx"
-    #     model = YOLO(Path(model_path, model_name), task="detect")
-    #     device_name = device("cpu")
-    #     IMGSIZE = 640 # required to make CPU inference a bit fast 
 
     predictions = []
     true_labels = []
     pred_boxes = []
     true_boxes_list = []  # List of lists, each inner list contains boxes for one image
     
+    print(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()
@@ -135,7 +126,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=False, imgsz=IMGSIZE)[0]
+        results = model.predict(example["image"], device=device_name, conf=THRESHOLD, verbose=True, imgsz=IMGSIZE)[0]
         pred_has_smoke = len(results) > 0
         predictions.append(int(pred_has_smoke))