Files

57 lines
1.7 KiB
Python
Raw Permalink Normal View History

import uuid
import cv2
import numpy as np
from fastapi import APIRouter, File, Form, HTTPException, UploadFile
from engine.loader import get_engine, get_inference_lock
from engine.annotator import annotate, to_base64
from schema import DefectBoxSchema, DetectResponse
router = APIRouter()
@router.post("/detect", response_model=DetectResponse)
async def detect(
image: UploadFile = File(...),
conf: float = Form(0.45),
return_annotated: bool = Form(False),
):
raw = await image.read()
arr = np.frombuffer(raw, dtype=np.uint8)
img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
if img is None:
raise HTTPException(status_code=400, detail="无法解码图像,请检查文件格式")
engine = get_engine()
async with get_inference_lock():
result = engine.detect(img, conf=conf)
annotated_b64 = None
if return_annotated:
annotated_b64 = to_base64(annotate(img, result))
return DetectResponse(
task_id=str(uuid.uuid4()),
verdict=result.verdict,
defect_count=result.defect_count,
max_severity=result.max_severity,
avg_confidence=result.avg_confidence,
duration_ms=result.duration_ms,
image_width=result.image_width,
image_height=result.image_height,
model_version=result.model_version,
defects=[
DefectBoxSchema(
class_id=d.class_id,
class_name=d.class_name,
class_name_zh=d.class_name_zh,
confidence=d.confidence,
severity=d.severity,
box_xyxy=d.box_xyxy,
)
for d in result.defects
],
annotated_image_base64=annotated_b64,
)