Spaces:
Sleeping
Sleeping
File size: 1,404 Bytes
bd8c54b cb2008f bd8c54b cb2008f 3b05912 cb2008f bd8c54b cb2008f 3b05912 bd8c54b cb2008f bd8c54b cb2008f bd8c54b cb2008f bd8c54b cb2008f bd8c54b 22d6743 3b05912 cb2008f 3b05912 bd8c54b 3b05912 22d6743 cb2008f 22d6743 cb2008f bd8c54b cb2008f 3b05912 22d6743 3b05912 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
from fastapi import FastAPI, File, UploadFile
from ultralytics import YOLO
import shutil
import uuid
from pathlib import Path
import base64
# Init FastAPI
app = FastAPI()
# Load YOLO model (place best.pt in /app folder)
model = YOLO("best.pt")
@app.get("/")
async def root():
return {"status": "ok", "message": "YOLOv8 Yellow Disease Detection API is running 🚀"}
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
temp_file = Path(f"/tmp/{uuid.uuid4()}_{file.filename}")
with temp_file.open("wb") as buffer:
shutil.copyfileobj(file.file, buffer)
results = model.predict(source=str(temp_file), save=False, conf=0.5)
detections = []
disease_found = False
for r in results:
for box in r.boxes:
cls_name = model.names[int(box.cls)]
conf = float(box.conf)
detections.append({
"class": cls_name,
"confidence": conf,
"bbox": [float(v) for v in box.xyxy[0].tolist()]
})
if cls_name == "yellow_disease":
disease_found = True
temp_file.unlink(missing_ok=True)
return {
"spray_signal": 1 if disease_found else 0, # ESC reads this
"frame_classification": "yellow_disease" if disease_found else "healthy",
"detections": detections # keep for logs/debugging
}
|