ilaandClaude Sonnet 4.6
68e4f444c5
- 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 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
|