69 lines
2.0 KiB
Python
69 lines
2.0 KiB
Python
import asyncio
|
|
import io
|
|
import os
|
|
from typing import List
|
|
|
|
from fastapi import FastAPI, File, HTTPException, UploadFile
|
|
from PIL import Image, UnidentifiedImageError
|
|
|
|
app = FastAPI(title="AI Plant Detection Service")
|
|
|
|
_model = None
|
|
MODEL_CACHE_DIR = os.environ.get("MODEL_CACHE_DIR", "/models")
|
|
MAX_FILE_SIZE = 10 * 1024 * 1024 # 10 MB
|
|
|
|
|
|
def get_model():
|
|
global _model
|
|
if _model is None:
|
|
from ultralytics import YOLO
|
|
os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
|
|
try:
|
|
_model = YOLO("foduucom/plant-leaf-detection-and-classification")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Impossible de charger le modèle YOLO: {e}") from e
|
|
return _model
|
|
|
|
|
|
def _run_inference(img: Image.Image) -> list:
|
|
model = get_model()
|
|
results = model.predict(img, conf=0.25, iou=0.45, verbose=False)
|
|
detections = []
|
|
if results and results[0].boxes:
|
|
boxes = results[0].boxes
|
|
names = model.names
|
|
for i in range(min(3, len(boxes))):
|
|
cls_id = int(boxes.cls[i].item())
|
|
conf = float(boxes.conf[i].item())
|
|
detections.append({
|
|
"class_name": names[cls_id],
|
|
"confidence": round(conf, 3),
|
|
})
|
|
return detections
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok", "model_loaded": _model is not None}
|
|
|
|
|
|
@app.post("/detect")
|
|
async def detect(file: UploadFile = File(...)):
|
|
data = await file.read()
|
|
|
|
if len(data) > MAX_FILE_SIZE:
|
|
raise HTTPException(status_code=413, detail="Fichier trop volumineux (max 10 MB)")
|
|
|
|
try:
|
|
img = Image.open(io.BytesIO(data)).convert("RGB")
|
|
except (UnidentifiedImageError, OSError) as e:
|
|
raise HTTPException(status_code=400, detail=f"Image invalide: {e}")
|
|
|
|
try:
|
|
loop = asyncio.get_event_loop()
|
|
detections = await loop.run_in_executor(None, _run_inference, img)
|
|
except RuntimeError as e:
|
|
raise HTTPException(status_code=503, detail=str(e))
|
|
|
|
return detections
|