diff --git a/ai-service/app/main.py b/ai-service/app/main.py index c7890af..260dc61 100644 --- a/ai-service/app/main.py +++ b/ai-service/app/main.py @@ -1,9 +1,30 @@ -from fastapi import FastAPI, File, HTTPException, UploadFile +import shutil +import subprocess +import threading +import time + +from fastapi import FastAPI, File, HTTPException, Request, UploadFile from . import config from .detector import MockDetector, ONNXDetector app = FastAPI(title="Silk AI Service", version="0.1.0") +START_TIME = time.time() +_lock = threading.Lock() +_requests = 0 +_latency_total = 0.0 + + +@app.middleware("http") +async def count_requests(request: Request, call_next): + global _requests, _latency_total + start = time.time() + response = await call_next(request) + latency = (time.time() - start) * 1000 + with _lock: + _requests += 1 + _latency_total += latency + return response if config.MODEL_MODE == "onnx": detector = ONNXDetector(config.MODEL_PATH, tuple(config.MODEL_LABELS)) @@ -36,3 +57,66 @@ async def detect(file: UploadFile = File(...)): for d in detections ], } + + +try: + import cv2 # noqa: E402 +except ImportError: + cv2 = None + + +@app.post("/stream-detect") +async def stream_detect(body: dict): + """摄像头流 AI 巡检骨架(#23):拉流抽帧 → 检测。 + + 二期功能:正式接入前需补摄像头视角数据与巡检任务编排。 + """ + url = (body or {}).get("url") or "" + if not url: + raise HTTPException(status_code=400, detail="缺少 url(rtsp/http-flv 流地址)") + if cv2 is None: + raise HTTPException(status_code=503, detail="OpenCV 未安装,无法拉流") + capture = cv2.VideoCapture(url) + if not capture.isOpened(): + raise HTTPException(status_code=502, detail="无法连接视频流") + try: + frames = 0 + detections = [] + while frames < 3: + ok, frame = capture.read() + if not ok: + break + ok_encode, buf = cv2.imencode(".jpg", frame) + if ok_encode: + detections.extend(detector.detect(buf.tobytes())) + frames += 1 + finally: + capture.release() + return {"model": config.MODEL_MODE, "frames": frames, "detections": detections[:10]} + + +@app.get("/metrics") +def metrics(): + """AI 服务监控(#26 骨架):请求量/延迟/GPU 信息(nvidia-smi 可选)。""" + with _lock: + reqs = _requests + avg = _latency_total / _requests if _requests else 0.0 + gpu = None + if shutil.which("nvidia-smi"): + try: + out = subprocess.run( + ["nvidia-smi", "--query-gpu=name,temperature.gpu,memory.used,memory.total", "--format=csv,noheader,nounits"], + capture_output=True, text=True, timeout=10, + ) + if out.returncode == 0 and out.stdout.strip(): + parts = [p.strip() for p in out.stdout.strip().split(",")] + gpu = {"name": parts[0], "tempC": parts[1], "memUsedMB": parts[2], "memTotalMB": parts[3]} + except Exception: + gpu = None + return { + "model": config.MODEL_MODE, + "uptimeSeconds": int(time.time() - START_TIME), + "requests": reqs, + "avgLatencyMs": round(avg, 2), + "gpu": gpu, + } diff --git a/ai-service/requirements.txt b/ai-service/requirements.txt index c396cdf..3994f12 100644 --- a/ai-service/requirements.txt +++ b/ai-service/requirements.txt @@ -7,3 +7,4 @@ pydantic>=2.5 python-multipart>=0.0.9 pytest>=8.0 httpx>=0.27 +opencv-python-headless>=4.8 diff --git a/ai-service/tests/test_api.py b/ai-service/tests/test_api.py index f5b0b35..9b8e7c7 100644 --- a/ai-service/tests/test_api.py +++ b/ai-service/tests/test_api.py @@ -40,3 +40,16 @@ def test_detect_empty_file_rejected(): def test_detect_invalid_image_rejected(): r = client.post("/detect", files={"file": ("a.png", b"junk", "image/png")}) assert r.status_code == 400 + + +def test_stream_detect_missing_url(): + r = client.post("/stream-detect", json={}) + assert r.status_code == 400 + + +def test_metrics_shape(): + r = client.get("/metrics") + assert r.status_code == 200 + body = r.json() + for key in ("model", "uptimeSeconds", "requests", "avgLatencyMs", "gpu"): + assert key in body