Update tasks/image.py

tasks/image.py

@@ -96,8 +96,8 @@ async def evaluate_image(request: ImageEvaluationRequest):
     test_dataset = train_test["test"]
 
     model = YOLOv10("best.pt")
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = model.to(device)
+    #device = "cuda" if torch.cuda.is_available() else "cpu"
+    #model = model.to(device)
     
     # Start tracking emissions
     tracker.start()
@@ -122,7 +122,7 @@ async def evaluate_image(request: ImageEvaluationRequest):
         # Make prediction with model
         image = example['image']
         with torch.inference_mode():
-            pred = model(image, conf=0.03)[0]
+            pred = model(image, conf=0.03, imgsz=1280)[0]
         smoke_detected = pred.boxes.xywhn.shape[0] > 0
         predictions.append(int(smoke_detected))