48 lines
1.2 KiB
Python
48 lines
1.2 KiB
Python
import io
|
|
import os
|
|
from typing import List
|
|
|
|
from fastapi import FastAPI, File, UploadFile
|
|
from PIL import Image
|
|
|
|
app = FastAPI(title="AI Plant Detection Service")
|
|
|
|
_model = None
|
|
MODEL_CACHE_DIR = os.environ.get("MODEL_CACHE_DIR", "/models")
|
|
|
|
|
|
def get_model():
|
|
global _model
|
|
if _model is None:
|
|
from ultralytics import YOLO
|
|
os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
|
|
_model = YOLO("foduucom/plant-leaf-detection-and-classification")
|
|
return _model
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.post("/detect")
|
|
async def detect(file: UploadFile = File(...)):
|
|
data = await file.read()
|
|
img = Image.open(io.BytesIO(data)).convert("RGB")
|
|
|
|
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
|