diff --git a/.gitignore b/.gitignore index eed0e79..4c11052 100644 --- a/.gitignore +++ b/.gitignore @@ -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/ diff --git a/ai-service/app/__init__.py b/ai-service/app/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/ai-service/app/__init__.py @@ -0,0 +1 @@ + diff --git a/ai-service/app/config.py b/ai-service/app/config.py new file mode 100644 index 0000000..95593a0 --- /dev/null +++ b/ai-service/app/config.py @@ -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")) diff --git a/ai-service/app/detector.py b/ai-service/app/detector.py new file mode 100644 index 0000000..8095515 --- /dev/null +++ b/ai-service/app/detector.py @@ -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 diff --git a/ai-service/app/main.py b/ai-service/app/main.py new file mode 100644 index 0000000..c7890af --- /dev/null +++ b/ai-service/app/main.py @@ -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 + ], + } diff --git a/ai-service/app/schemas.py b/ai-service/app/schemas.py new file mode 100644 index 0000000..0f1921e --- /dev/null +++ b/ai-service/app/schemas.py @@ -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 diff --git a/ai-service/tests/__init__.py b/ai-service/tests/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/ai-service/tests/__init__.py @@ -0,0 +1 @@ + diff --git a/ai-service/tests/test_api.py b/ai-service/tests/test_api.py new file mode 100644 index 0000000..f5b0b35 --- /dev/null +++ b/ai-service/tests/test_api.py @@ -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 diff --git a/ai-service/tests/test_detector.py b/ai-service/tests/test_detector.py new file mode 100644 index 0000000..756203a --- /dev/null +++ b/ai-service/tests/test_detector.py @@ -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")