weijuesen
123 lines
3.7 KiB
Python
123 lines
3.7 KiB
Python
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))
|
||
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
|
||
],
|
||
}
|
||
|
||
|
||
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,
|
||
}
|