- engine/defect_classes: HRIPCB 6类缺陷定义 + 4类扩展 + severity排序 - engine/base: DefectBox / DetectResult 数据类 + BaseInferenceEngine 抽象基类 - engine/ultralytics_adapter: YOLOv8 .pt 推理适配器 + warmup - engine/annotator: 按 severity 着色的标注图生成 + base64 编码 - engine/loader: 单例管理 + asyncio.Lock 推理串行化 + 热重载支持 - config: Pydantic v1 BaseSettings,读取 .env - schema: DetectResponse / HealthResponse / ReloadRequest 等响应模型 - api/detect: POST /api/detect,内存读图,推理锁保护 - api/health: GET /api/health - api/model: POST /api/model/reload 热重载接口 - main: FastAPI lifespan 启动加载模型 测试通过:/api/health 200,/api/detect 空图返回 pass,标注图 base64 正常 Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
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,
|
|
)
|