app.py

from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from ultralytics import YOLO
from PIL import Image
import io

app = FastAPI()

# Load the YOLOv8 model
model = YOLO("best.pt")

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        image = Image.open(io.BytesIO(contents)).convert("RGB")

        results = model.predict(image)
        prediction = results[0]

        if len(prediction.boxes.cls) == 0:
            return JSONResponse({
                "status": "Healthy",
                "message": "No disease detected",
                "predictions": []
            })

        response = {
            "status": "Diseased",
            "predictions": []
        }

        for box in prediction.boxes:
            class_id = int(box.cls[0])
            confidence = float(box.conf[0])
            label = model.names[class_id]

            response["predictions"].append({
                "disease": label,
                "confidence": round(confidence, 3),
                "status": "Diseased" if confidence > 0.5 else "Uncertain"
            })

        return JSONResponse(response)

    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})