chore: 允许 ai-service Python 文件入库（.gitignore 例外）

2026-08-12 16:22:45 +08:00
parent 99137a6fb4
commit b89ad5e865
9 changed files with 278 additions and 0 deletions
@@ -28,6 +28,13 @@ __pycache__/
 *.ps1
 *.py
 *.sql
+# ai-service（Python 推理服务）文件例外
+!ai-service/
+!ai-service/app/
+!ai-service/app/*.py
+!ai-service/tests/
+!ai-service/tests/*.py
+.venv/
 #故障排查处理记录（微信小程序）.md
 app/android/.gradle/
 app/android/build/
@@ -0,0 +1 @@
+
@@ -0,0 +1,10 @@
+import os
+
+# 运行模式：mock（默认，无模型文件）| onnx（加载 best.onnx 真实推理）
+MODEL_MODE = os.getenv("MODEL_MODE", "mock")
+MODEL_PATH = os.getenv("MODEL_PATH", "models/best.onnx")
+# YOLO 类别（二分类训练基线：healthy/sick；7 类病种扩展后再调整）
+MODEL_LABELS = os.getenv("MODEL_LABELS", "healthy,sick").split(",")
+# mock 模式返回的固定结果
+MOCK_CLASS = os.getenv("MOCK_CLASS", "healthy")
+MOCK_CONFIDENCE = float(os.getenv("MOCK_CONFIDENCE", "0.95"))
@@ -0,0 +1,140 @@
+import io
+import os
+from abc import ABC, abstractmethod
+
+from PIL import Image
+
+
+class Detector(ABC):
+    @abstractmethod
+    def detect(self, image_bytes: bytes) -> list[dict]:
+        """返回 [{"bbox": {x,y,w,h}, "class_name", "confidence"}]"""
+
+
+class MockDetector(Detector):
+    """开发联调用：校验图片可解析后返回固定检测结果。"""
+
+    def __init__(self, class_name: str = "healthy", confidence: float = 0.95):
+        self.class_name = class_name
+        self.confidence = confidence
+
+    def detect(self, image_bytes: bytes) -> list[dict]:
+        try:
+            with Image.open(io.BytesIO(image_bytes)) as img:
+                img.verify()
+        except Exception as exc:
+            raise ValueError("无法解析图片") from exc
+        return [
+            {
+                "bbox": {"x": 10, "y": 10, "w": 100, "h": 100},
+                "class_name": self.class_name,
+                "confidence": self.confidence,
+            }
+        ]
+
+
+class ONNXDetector(Detector):
+    """YOLOv8 导出的 best.onnx 推理（CPU）。
+
+    YOLO 训练恢复后放置模型文件，MODEL_MODE=onnx 生效。
+    输出张量按 YOLOv8 常见格式 [1, 4+nc, 8400]（或转置）解析，含 NMS。
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        labels: tuple[str, ...],
+        input_size: int = 640,
+        conf_threshold: float = 0.25,
+        iou_threshold: float = 0.45,
+    ):
+        import numpy as np
+        import onnxruntime as ort
+
+        if not os.path.exists(model_path):
+            raise FileNotFoundError(f"模型文件不存在: {model_path}")
+        self.session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
+        self.labels = labels
+        self.input_size = input_size
+        self.conf_threshold = conf_threshold
+        self.iou_threshold = iou_threshold
+        self.np = np
+
+    def detect(self, image_bytes: bytes) -> list[dict]:
+        np = self.np
+        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        orig_w, orig_h = img.size
+        scale = min(self.input_size / orig_w, self.input_size / orig_h)
+        new_w, new_h = int(round(orig_w * scale)), int(round(orig_h * scale))
+        resized = img.resize((new_w, new_h))
+        canvas = Image.new("RGB", (self.input_size, self.input_size), (114, 114, 114))
+        pad_x = (self.input_size - new_w) // 2
+        pad_y = (self.input_size - new_h) // 2
+        canvas.paste(resized, (pad_x, pad_y))
+        blob = np.asarray(canvas, dtype=np.float32) / 255.0
+        blob = blob.transpose(2, 0, 1)[None, ...]
+        out = self.session.run(None, {self.session.get_inputs()[0].name: blob})[0]
+
+        nc = len(self.labels)
+        if out.shape[1] != 4 + nc:
+            out = out.transpose(0, 2, 1)
+        out = out[0]
+        if out.shape[0] == 4 + nc:
+            out = out.T
+
+        boxes, scores, class_ids = [], [], []
+        for row in out:
+            cx, cy, w, h = row[:4]
+            cls_scores = row[4 : 4 + nc]
+            cls_id = int(cls_scores.argmax())
+            score = float(cls_scores[cls_id])
+            if score < self.conf_threshold:
+                continue
+            boxes.append(
+                [
+                    (cx - w / 2 - pad_x) / scale,
+                    (cy - h / 2 - pad_y) / scale,
+                    w / scale,
+                    h / scale,
+                ]
+            )
+            scores.append(score)
+            class_ids.append(cls_id)
+
+        keep = self._nms(np.asarray(boxes), np.asarray(scores))
+        return [
+            {
+                "bbox": {
+                    "x": float(boxes[i][0]),
+                    "y": float(boxes[i][1]),
+                    "w": float(boxes[i][2]),
+                    "h": float(boxes[i][3]),
+                },
+                "class_name": self.labels[class_ids[i]],
+                "confidence": float(scores[i]),
+            }
+            for i in keep
+        ]
+
+    def _nms(self, boxes, scores):
+        np = self.np
+        if len(boxes) == 0:
+            return []
+        x1 = boxes[:, 0]
+        y1 = boxes[:, 1]
+        x2 = boxes[:, 0] + boxes[:, 2]
+        y2 = boxes[:, 1] + boxes[:, 3]
+        areas = boxes[:, 2] * boxes[:, 3]
+        order = scores.argsort()[::-1]
+        keep = []
+        while order.size > 0:
+            i = int(order[0])
+            keep.append(i)
+            xx1 = np.maximum(x1[i], x1[order[1:]])
+            yy1 = np.maximum(y1[i], y1[order[1:]])
+            xx2 = np.minimum(x2[i], x2[order[1:]])
+            yy2 = np.minimum(y2[i], y2[order[1:]])
+            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
+            iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-6)
+            order = order[1:][iou <= self.iou_threshold]
+        return keep
@@ -0,0 +1,38 @@
+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
+        ],
+    }
@@ -0,0 +1,14 @@
+from pydantic import BaseModel
+
+
+class BBox(BaseModel):
+    x: float
+    y: float
+    w: float
+    h: float
+
+
+class Detection(BaseModel):
+    bbox: BBox
+    class_name: str
+    confidence: float
@@ -0,0 +1 @@
+
@@ -0,0 +1,42 @@
+import base64
+
+from fastapi.testclient import TestClient
+
+from app.main import app
+
+# 1x1 透明 PNG
+TINY_PNG = base64.b64decode(
+    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
+)
+
+client = TestClient(app)
+
+
+def test_health():
+    r = client.get("/health")
+    assert r.status_code == 200
+    body = r.json()
+    assert body["status"] == "ok"
+    assert body["model"] in ("mock", "onnx")
+
+
+def test_detect_ok():
+    r = client.post("/detect", files={"file": ("a.png", TINY_PNG, "image/png")})
+    assert r.status_code == 200
+    body = r.json()
+    assert body["model"] == "mock"
+    assert len(body["detections"]) >= 1
+    d = body["detections"][0]
+    assert d["class"] in ("healthy", "sick")
+    assert 0 <= d["confidence"] <= 1
+    assert d["bbox"]["w"] > 0
+
+
+def test_detect_empty_file_rejected():
+    r = client.post("/detect", files={"file": ("a.png", b"", "image/png")})
+    assert r.status_code == 400
+
+
+def test_detect_invalid_image_rejected():
+    r = client.post("/detect", files={"file": ("a.png", b"junk", "image/png")})
+    assert r.status_code == 400
@@ -0,0 +1,25 @@
+import base64
+
+import pytest
+
+from app.detector import MockDetector
+
+# 1x1 透明 PNG
+TINY_PNG = base64.b64decode(
+    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
+)
+
+
+def test_mock_detector_returns_canned_detection():
+    det = MockDetector(class_name="sick", confidence=0.92)
+    result = det.detect(TINY_PNG)
+    assert len(result) == 1
+    assert result[0]["class_name"] == "sick"
+    assert result[0]["confidence"] == 0.92
+    assert result[0]["bbox"]["w"] > 0
+
+
+def test_mock_detector_rejects_invalid_image():
+    det = MockDetector()
+    with pytest.raises(ValueError):
+        det.detect(b"not an image")