| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859 |
- 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)
|
