weijuesen
39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
from fastapi import FastAPI, File, HTTPException, UploadFile
|
|
|
|
from . import config
|
|
from .detector import MockDetector, ONNXDetector
|
|
|
|
app = FastAPI(title="Silk AI Service", version="0.1.0")
|
|
|
|
if config.MODEL_MODE == "onnx":
|
|
detector = ONNXDetector(config.MODEL_PATH, tuple(config.MODEL_LABELS))
|
|
else:
|
|
detector = MockDetector(class_name=config.MOCK_CLASS, confidence=config.MOCK_CONFIDENCE)
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok", "model": config.MODEL_MODE}
|
|
|
|
|
|
@app.post("/detect")
|
|
async def detect(file: UploadFile = File(...)):
|
|
data = await file.read()
|
|
if not data:
|
|
raise HTTPException(status_code=400, detail="图片为空")
|
|
try:
|
|
detections = detector.detect(data)
|
|
except ValueError as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
return {
|
|
"model": config.MODEL_MODE,
|
|
"detections": [
|
|
{
|
|
"bbox": d["bbox"],
|
|
"class": d["class_name"],
|
|
"confidence": d["confidence"],
|
|
}
|
|
for d in detections
|
|
],
|
|
}
|