Update tasks/image.py

tasks/image.py

@@ -180,4 +180,41 @@ async def evaluate_image(request: ImageEvaluationRequest):
         }
     }
     
-    return results 
+    return results 
+
+@router.post("/detect-smoke", tags=["Smoke Detection"])
+async def detect_smoke_image(file: UploadFile = File(...)):
+    """
+    Detect smoke in a single uploaded image using YOLO model.
+    Returns whether smoke is detected and any bounding boxes.
+    """
+    # Load the uploaded image
+    image_bytes = await file.read()
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    
+    # Load model
+    model_path = Path("tasks", "models")
+    model_name = "best.pt"
+    model = YOLO(Path(model_path, model_name))
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+    
+    # Run prediction
+    results = model.predict(image, verbose=False)
+    boxes = results[0].boxes
+
+    response = {
+        "smoke_detected": bool(len(boxes)),
+        "num_detections": len(boxes),
+        "boxes": []
+    }
+
+    for box in boxes:
+        xywhn = box.xywhn[0].cpu().numpy().tolist()
+        conf = float(box.conf[0]) if box.conf is not None else None
+        response["boxes"].append({
+            "xywhn": xywhn,
+            "confidence": conf
+        })
+
+    return JSONResponse(content=response)