feat(vision): 实现 mingxi-vision 推理服务核心功能
- 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>
This commit is contained in:
@@ -0,0 +1,4 @@
|
|||||||
|
MODEL_PATH=./models/best.pt
|
||||||
|
RUNTIME=ultralytics
|
||||||
|
DEVICE=cuda
|
||||||
|
CONF_THRESHOLD=0.45
|
||||||
@@ -0,0 +1,56 @@
|
|||||||
|
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,
|
||||||
|
)
|
||||||
|
|||||||
@@ -0,0 +1,18 @@
|
|||||||
|
from fastapi import APIRouter
|
||||||
|
|
||||||
|
from engine.loader import get_engine
|
||||||
|
from config import settings
|
||||||
|
from schema import HealthResponse
|
||||||
|
|
||||||
|
router = APIRouter()
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/health", response_model=HealthResponse)
|
||||||
|
async def health():
|
||||||
|
engine = get_engine()
|
||||||
|
return HealthResponse(
|
||||||
|
status="ok",
|
||||||
|
model_version=engine.model_version,
|
||||||
|
runtime=settings.runtime,
|
||||||
|
device=settings.device,
|
||||||
|
)
|
||||||
|
|||||||
@@ -0,0 +1,16 @@
|
|||||||
|
from fastapi import APIRouter, HTTPException
|
||||||
|
|
||||||
|
from engine.loader import get_inference_lock, reload_engine
|
||||||
|
from schema import ReloadRequest, ReloadResponse
|
||||||
|
|
||||||
|
router = APIRouter()
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/model/reload", response_model=ReloadResponse)
|
||||||
|
async def model_reload(req: ReloadRequest):
|
||||||
|
async with get_inference_lock():
|
||||||
|
try:
|
||||||
|
engine = reload_engine(req.model_path, req.runtime)
|
||||||
|
except Exception as e:
|
||||||
|
raise HTTPException(status_code=500, detail=f"模型重载失败: {e}")
|
||||||
|
return ReloadResponse(status="ok", model_version=engine.model_version)
|
||||||
|
|||||||
@@ -0,0 +1,16 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
from pydantic import BaseSettings # pydantic v1
|
||||||
|
|
||||||
|
|
||||||
|
class Settings(BaseSettings):
|
||||||
|
model_path: str = str(Path(__file__).parent / "models" / "best.pt")
|
||||||
|
runtime: str = "ultralytics" # ultralytics | onnx
|
||||||
|
device: str = "cuda" # cuda | cpu
|
||||||
|
conf_threshold: float = 0.45
|
||||||
|
|
||||||
|
class Config:
|
||||||
|
env_file = ".env"
|
||||||
|
env_file_encoding = "utf-8"
|
||||||
|
|
||||||
|
|
||||||
|
settings = Settings()
|
||||||
|
|||||||
@@ -0,0 +1,37 @@
|
|||||||
|
import base64
|
||||||
|
from typing import Tuple
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from .base import DetectResult
|
||||||
|
|
||||||
|
# 按严重程度着色,BGR格式
|
||||||
|
SEVERITY_COLORS: dict = {
|
||||||
|
"fatal": (0, 0, 220),
|
||||||
|
"major": (0, 128, 255),
|
||||||
|
"minor": (0, 215, 255),
|
||||||
|
"rework": (255, 165, 0),
|
||||||
|
"none": (180, 180, 180),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def annotate(image: np.ndarray, result: DetectResult) -> np.ndarray:
|
||||||
|
img = image.copy()
|
||||||
|
for d in result.defects:
|
||||||
|
color: Tuple = SEVERITY_COLORS.get(d.severity, (180, 180, 180))
|
||||||
|
x1, y1, x2, y2 = [int(v) for v in d.box_xyxy]
|
||||||
|
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
|
||||||
|
label = f"{d.class_name_zh} {d.confidence:.2f}"
|
||||||
|
(tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
|
||||||
|
cv2.rectangle(img, (x1, y1 - th - 6), (x1 + tw + 4, y1), color, -1)
|
||||||
|
cv2.putText(img, label, (x1 + 2, y1 - 4),
|
||||||
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
|
||||||
|
return img
|
||||||
|
|
||||||
|
|
||||||
|
def to_base64(image: np.ndarray, quality: int = 85) -> str:
|
||||||
|
ok, buf = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, quality])
|
||||||
|
if not ok:
|
||||||
|
raise RuntimeError("图像编码失败")
|
||||||
|
return base64.b64encode(buf.tobytes()).decode()
|
||||||
|
|||||||
@@ -0,0 +1,59 @@
|
|||||||
|
from abc import ABC, abstractmethod
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from .defect_classes import max_severity
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DefectBox:
|
||||||
|
class_id: int
|
||||||
|
class_name: str
|
||||||
|
class_name_zh: str
|
||||||
|
confidence: float
|
||||||
|
severity: str # fatal / major / minor / rework / none
|
||||||
|
box_xyxy: List[float]
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DetectResult:
|
||||||
|
defects: List[DefectBox]
|
||||||
|
duration_ms: float
|
||||||
|
image_width: int
|
||||||
|
image_height: int
|
||||||
|
model_version: str
|
||||||
|
|
||||||
|
@property
|
||||||
|
def defect_count(self) -> int:
|
||||||
|
return len(self.defects)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def max_severity(self) -> str:
|
||||||
|
return max_severity([d.severity for d in self.defects]) if self.defects else "none"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def avg_confidence(self) -> float:
|
||||||
|
if not self.defects:
|
||||||
|
return 0.0
|
||||||
|
return round(sum(d.confidence for d in self.defects) / len(self.defects), 4)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def verdict(self) -> str:
|
||||||
|
return "fail" if self.defects else "pass"
|
||||||
|
|
||||||
|
|
||||||
|
class BaseInferenceEngine(ABC):
|
||||||
|
@abstractmethod
|
||||||
|
def detect(self, image: np.ndarray, conf: float = 0.45) -> DetectResult:
|
||||||
|
...
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def warmup(self) -> None:
|
||||||
|
...
|
||||||
|
|
||||||
|
@property
|
||||||
|
@abstractmethod
|
||||||
|
def model_version(self) -> str:
|
||||||
|
...
|
||||||
|
|||||||
@@ -0,0 +1,33 @@
|
|||||||
|
from typing import Dict
|
||||||
|
|
||||||
|
DEFECT_CLASSES: Dict[int, dict] = {
|
||||||
|
# HRIPCB 数据集原始 6 类
|
||||||
|
0: {"name": "missing_hole", "zh": "缺孔", "severity": "fatal"},
|
||||||
|
1: {"name": "mouse_bite", "zh": "鼠咬", "severity": "major"},
|
||||||
|
2: {"name": "open_circuit", "zh": "断路", "severity": "fatal"},
|
||||||
|
3: {"name": "short_circuit", "zh": "短路", "severity": "fatal"},
|
||||||
|
4: {"name": "spur", "zh": "毛刺", "severity": "minor"},
|
||||||
|
5: {"name": "spurious_copper", "zh": "余铜", "severity": "major"},
|
||||||
|
# 扩展类(Phase 2 标注后加入训练)
|
||||||
|
6: {"name": "oxidation", "zh": "氧化", "severity": "minor"},
|
||||||
|
7: {"name": "solder_ball", "zh": "锡珠", "severity": "rework"},
|
||||||
|
8: {"name": "scratch", "zh": "划痕", "severity": "minor"},
|
||||||
|
9: {"name": "label_error", "zh": "标签错贴","severity": "rework"},
|
||||||
|
}
|
||||||
|
|
||||||
|
# 严重程度从低到高,用于比较 max_severity
|
||||||
|
SEVERITY_ORDER = ["none", "rework", "minor", "major", "fatal"]
|
||||||
|
|
||||||
|
|
||||||
|
def get_class_info(class_id: int) -> dict:
|
||||||
|
return DEFECT_CLASSES.get(class_id, {
|
||||||
|
"name": f"unknown_{class_id}",
|
||||||
|
"zh": f"未知_{class_id}",
|
||||||
|
"severity": "minor",
|
||||||
|
})
|
||||||
|
|
||||||
|
|
||||||
|
def max_severity(severities) -> str:
|
||||||
|
if not severities:
|
||||||
|
return "none"
|
||||||
|
return max(severities, key=lambda s: SEVERITY_ORDER.index(s) if s in SEVERITY_ORDER else 0)
|
||||||
|
|||||||
@@ -0,0 +1,43 @@
|
|||||||
|
import asyncio
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from .base import BaseInferenceEngine
|
||||||
|
|
||||||
|
_engine: Optional[BaseInferenceEngine] = None
|
||||||
|
# 懒初始化,确保在事件循环启动后创建
|
||||||
|
_inference_lock: Optional[asyncio.Lock] = None
|
||||||
|
|
||||||
|
|
||||||
|
def get_engine() -> BaseInferenceEngine:
|
||||||
|
if _engine is None:
|
||||||
|
raise RuntimeError("推理引擎未初始化,请先调用 init_engine()")
|
||||||
|
return _engine
|
||||||
|
|
||||||
|
|
||||||
|
def get_inference_lock() -> asyncio.Lock:
|
||||||
|
global _inference_lock
|
||||||
|
if _inference_lock is None:
|
||||||
|
_inference_lock = asyncio.Lock()
|
||||||
|
return _inference_lock
|
||||||
|
|
||||||
|
|
||||||
|
def init_engine(model_path: str, runtime: str = "ultralytics") -> BaseInferenceEngine:
|
||||||
|
global _engine
|
||||||
|
if runtime == "onnx":
|
||||||
|
from .onnx_adapter import OnnxRuntimeAdapter
|
||||||
|
_engine = OnnxRuntimeAdapter(model_path)
|
||||||
|
else:
|
||||||
|
from .ultralytics_adapter import UltralyticsAdapter
|
||||||
|
_engine = UltralyticsAdapter(model_path)
|
||||||
|
|
||||||
|
print(f"[mingxi-vision] 加载模型: {model_path}")
|
||||||
|
_engine.warmup()
|
||||||
|
print(f"[mingxi-vision] 引擎就绪: {_engine.__class__.__name__} · {_engine.model_version}")
|
||||||
|
return _engine
|
||||||
|
|
||||||
|
|
||||||
|
def reload_engine(model_path: str, runtime: str = "ultralytics") -> BaseInferenceEngine:
|
||||||
|
"""热重载模型,调用方必须持有推理锁"""
|
||||||
|
global _engine
|
||||||
|
_engine = None # 释放旧引用,让 GC 回收显存
|
||||||
|
return init_engine(model_path, runtime)
|
||||||
|
|||||||
@@ -0,0 +1,56 @@
|
|||||||
|
import time
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from .base import BaseInferenceEngine, DefectBox, DetectResult
|
||||||
|
from .defect_classes import get_class_info
|
||||||
|
|
||||||
|
|
||||||
|
class UltralyticsAdapter(BaseInferenceEngine):
|
||||||
|
|
||||||
|
def __init__(self, model_path: str):
|
||||||
|
from ultralytics import YOLO
|
||||||
|
self._model_path = str(model_path)
|
||||||
|
self._model = YOLO(self._model_path)
|
||||||
|
self._version = Path(model_path).stem
|
||||||
|
|
||||||
|
def detect(self, image: np.ndarray, conf: float = 0.45) -> DetectResult:
|
||||||
|
h, w = image.shape[:2]
|
||||||
|
t0 = time.perf_counter()
|
||||||
|
results = self._model(image, conf=conf, verbose=False)
|
||||||
|
duration_ms = (time.perf_counter() - t0) * 1000
|
||||||
|
|
||||||
|
defects = []
|
||||||
|
for r in results:
|
||||||
|
if r.boxes is None:
|
||||||
|
continue
|
||||||
|
for box in r.boxes:
|
||||||
|
class_id = int(box.cls[0])
|
||||||
|
confidence = float(box.conf[0])
|
||||||
|
xyxy = box.xyxy[0].tolist()
|
||||||
|
info = get_class_info(class_id)
|
||||||
|
defects.append(DefectBox(
|
||||||
|
class_id=class_id,
|
||||||
|
class_name=info["name"],
|
||||||
|
class_name_zh=info["zh"],
|
||||||
|
confidence=round(confidence, 4),
|
||||||
|
severity=info["severity"],
|
||||||
|
box_xyxy=[round(v, 1) for v in xyxy],
|
||||||
|
))
|
||||||
|
|
||||||
|
return DetectResult(
|
||||||
|
defects=defects,
|
||||||
|
duration_ms=round(duration_ms, 1),
|
||||||
|
image_width=w,
|
||||||
|
image_height=h,
|
||||||
|
model_version=self._version,
|
||||||
|
)
|
||||||
|
|
||||||
|
def warmup(self) -> None:
|
||||||
|
dummy = np.zeros((640, 640, 3), dtype=np.uint8)
|
||||||
|
self.detect(dummy, conf=0.45)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def model_version(self) -> str:
|
||||||
|
return self._version
|
||||||
|
|||||||
@@ -0,0 +1,26 @@
|
|||||||
|
from contextlib import asynccontextmanager
|
||||||
|
|
||||||
|
from fastapi import FastAPI
|
||||||
|
|
||||||
|
from config import settings
|
||||||
|
from engine.loader import init_engine
|
||||||
|
from api.detect import router as detect_router
|
||||||
|
from api.health import router as health_router
|
||||||
|
from api.model import router as model_router
|
||||||
|
|
||||||
|
|
||||||
|
@asynccontextmanager
|
||||||
|
async def lifespan(app: FastAPI):
|
||||||
|
init_engine(settings.model_path, settings.runtime)
|
||||||
|
yield
|
||||||
|
|
||||||
|
|
||||||
|
app = FastAPI(
|
||||||
|
title="明析推理服务",
|
||||||
|
version="1.0.0",
|
||||||
|
lifespan=lifespan,
|
||||||
|
)
|
||||||
|
|
||||||
|
app.include_router(detect_router, prefix="/api")
|
||||||
|
app.include_router(health_router, prefix="/api")
|
||||||
|
app.include_router(model_router, prefix="/api")
|
||||||
|
|||||||
@@ -0,0 +1,42 @@
|
|||||||
|
from typing import List, Optional
|
||||||
|
from pydantic import BaseModel # pydantic v1
|
||||||
|
|
||||||
|
|
||||||
|
class DefectBoxSchema(BaseModel):
|
||||||
|
class_id: int
|
||||||
|
class_name: str
|
||||||
|
class_name_zh: str
|
||||||
|
confidence: float
|
||||||
|
severity: str
|
||||||
|
box_xyxy: List[float]
|
||||||
|
|
||||||
|
|
||||||
|
class DetectResponse(BaseModel):
|
||||||
|
task_id: str
|
||||||
|
verdict: str # pass | fail
|
||||||
|
defect_count: int
|
||||||
|
max_severity: str
|
||||||
|
avg_confidence: float
|
||||||
|
duration_ms: float
|
||||||
|
image_width: int
|
||||||
|
image_height: int
|
||||||
|
model_version: str
|
||||||
|
defects: List[DefectBoxSchema]
|
||||||
|
annotated_image_base64: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
|
class ReloadRequest(BaseModel):
|
||||||
|
model_path: str
|
||||||
|
runtime: str = "ultralytics"
|
||||||
|
|
||||||
|
|
||||||
|
class ReloadResponse(BaseModel):
|
||||||
|
status: str
|
||||||
|
model_version: str
|
||||||
|
|
||||||
|
|
||||||
|
class HealthResponse(BaseModel):
|
||||||
|
status: str
|
||||||
|
model_version: str
|
||||||
|
runtime: str
|
||||||
|
device: str
|
||||||
|
|||||||
Reference in New Issue
Block a user