from fastapi import FastAPI, File, UploadFile, Body
from ultralytics import YOLO
import io
from PIL import Image

app = FastAPI()

# 1. Load your custom trained model
model = YOLO('best.pt') 

# 2. Global state for the confidence threshold
# Defaulting to 0.25 (YOLO's internal default)
current_conf = 0.25

@app.get("/get_confidence")
async def get_confidence():
    """Returns the current confidence threshold used by the model."""
    return {
        "status": "success",
        "current_confidence": current_conf,
        "model_version": "best.pt"
    }

@app.post("/set_confidence")
async def set_confidence(threshold: float = Body(..., embed=True)):
    """Updates the confidence threshold globally."""
    global current_conf
    if 0.0 <= threshold <= 1.0:
        current_conf = threshold
        return {"status": "success", "new_confidence": current_conf}
    else:
        return {"status": "error", "message": "Threshold must be between 0.0 and 1.0"}

@app.post("/detect")
async def detect_ripeness(file: UploadFile = File(...)):
    image_bytes = await file.read()
    img = Image.open(io.BytesIO(image_bytes))

    # 3. Apply the dynamic threshold to the inference
    results = model(img, conf=current_conf)

    detections = []
    for r in results:
        for box in r.boxes:
            detections.append({
                "class": model.names[int(box.cls)],
                "confidence": round(float(box.conf), 2),
                "box": box.xyxy.tolist()[0]
            })

    return {
        "status": "success", 
        "current_threshold": current_conf,
        "data": detections
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)