YOLO服务搭建

小赖同学啊2026-06-06 11:39

在水利工程 AI 运行管理平台中，搭建 YOLO 视频分析服务 是实现"人员违规入侵检测"、"水面漂浮物识别"、"未穿救生衣告警"等场景的核心。

为了保证与 Django 主业务系统解耦且具备高性能，推荐使用 FastAPI + YOLOv8 (Ultralytics) + OpenCV 的微服务架构。

以下是完整的 YOLO 服务搭建指南，从环境准备到代码实现，再到与 Django 的业务联动。

一、 技术选型与架构

  • AI 框架 : Ultralytics YOLOv8 (目前工业界落地最成熟，支持检测、分割，且 API 极其友好)
  • Web 框架 : FastAPI (原生异步，高并发，适合做推理 API)
  • 视频处理 : OpenCV (cv2) 用于拉取 RTSP 摄像头视频流并抽帧
  • 通信方式 : YOLO 服务检测到异常后，通过 HTTP Webhook 主动推送告警信息（含截图）给 Django 后端。

二、 环境准备

建议在带有 NVIDIA GPU 的服务器或开发机上运行，以保障实时推理速度。

bash 复制代码 
# 1. 创建虚拟环境
python -m venv yolo_env
source yolo_env/bin/activate  # Windows: yolo_env\Scripts\activate

# 2. 安装核心依赖
pip install fastapi uvicorn opencv-python-headless ultralytics requests pydantic python-multipart

# 3. (可选) 如果需要 GPU 加速，请确保已安装 CUDA 和 PyTorch GPU 版本
#
在水利工程 AI 运行管理平台中，搭建 **YOLO 视频分析服务** 是实现"人员违规入侵检测"、"水面漂浮物识别"、"未穿救生衣告警"等场景的核心。

为了保证与 Django 主业务系统解耦且具备高性能，推荐使用 **FastAPI + YOLOv8 (Ultralytics) + OpenCV** 的微服务架构。

以下是完整的 YOLO 服务搭建指南，从环境准备到代码实现，再到与 Django 的业务联动。

---

### 一、 技术选型与架构

*   **AI 框架**: `Ultralytics YOLOv8` (目前工业界落地最成熟，支持检测、分割，且 API 极其友好)
*   **Web 框架**: `FastAPI` (原生异步，高并发，适合做推理 API)
*   **视频处理**: `OpenCV` (`cv2`) 用于拉取 RTSP 摄像头视频流并抽帧
*   **通信方式**: YOLO 服务检测到异常后，通过 **HTTP Webhook** 主动推送告警信息（含截图）给 Django 后端。

---

### 二、 环境准备

建议在带有 NVIDIA GPU 的服务器或开发机上运行，以保障实时推理速度。

```bash
# 1. 创建虚拟环境
python -m venv yolo_env
source yolo_env/bin/activate  # Windows: yolo_env\Scripts\activate

# 2. 安装核心依赖
pip install fastapi uvicorn opencv-python-headless ultralytics requests pydantic python-multipart

# 3. (可选) 如果需要 GPU 加速，请确保已安装 CUDA 和 PyTorch GPU 版本
# pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

三、 核心代码实现

我们将服务拆分为两个文件：模型推理封装 (detector.py) 和 FastAPI 路由 (main.py)。

1. detector.py (YOLO 推理与流处理核心)
python 复制代码 
import cv2
import base64
import logging
from ultralytics import YOLO
from typing import List, Dict, Any

# 配置日志
logging.basicConfig(level=logging.INFO)

class WaterYOLODetector:
    def __init__(self, model_path: str = "yolov8n.pt", device: str = "0"):
        """
        初始化 YOLO 模型
        :param model_path: 模型权重路径 (可以是官方预训练模型，也可以是你自己微调的 .pt 文件)
        :param device: '0' 表示 GPU, 'cpu' 表示 CPU
        """
        logging.info(f"Loading YOLO model from {model_path} on device {device}...")
        self.model = YOLO(model_path)
        self.device = device
        
        # 水利场景关注的类别 ID (以 COCO 数据集为例，实际需根据你的自定义数据集调整)
        # 0: person (人员), 41: cup (漂浮物/垃圾), 等
        self.target_classes = {0: "person", 41: "cup"} 

    def process_frame(self, frame: cv2.Mat, conf_threshold: float = 0.5) -> Dict[str, Any]:
        """
        处理单帧图像，返回检测结果和标记后的图像
        """
        # 推理
        results = self.model(frame, device=self.device, conf=conf_threshold, verbose=False)
        
        detections = []
        annotated_frame = frame.copy()

        for result in results:
            boxes = result.boxes
            for box in boxes:
                cls_id = int(box.cls[0])
                if cls_id in self.target_classes:
                    class_name = self.target_classes[cls_id]
                    conf = float(box.conf[0])
                    x1, y1, x2, y2 = map(int, box.xyxy[0])
                    
                    detections.append({
                        "class": class_name,
                        "confidence": round(conf, 2),
                        "bbox": [x1, y1, x2, y2]
                    })
                    
                    # 在图像上画框 (B, G, R)
                    color = (0, 0, 255) if class_name == "person" else (0, 255, 0)
                    cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color, 2)
                    cv2.putText(annotated_frame, f"{class_name} {conf:.2f}", (x1, y1 - 10),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

        return {
            "detections": detections,
            "annotated_frame": annotated_frame
        }

    def encode_image_to_base64(self, frame: cv2.Mat) -> str:
        """将 OpenCV 图像编码为 Base64 字符串，便于通过 HTTP 传输"""
        _, buffer = cv2.imencode('.jpg', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70]) # 70% 质量以减小体积
        return base64.b64encode(buffer).decode('utf-8')
2. main.py (FastAPI 服务入口与业务联动)
python 复制代码 
from fastapi import FastAPI, BackgroundTasks, HTTPException
from pydantic import BaseModel
import cv2
import requests
import threading
import time
from detector import WaterYOLODetector

app = FastAPI(title="Water Conservancy YOLO AI Service")

# 全局初始化模型 (避免每次请求都加载模型)
detector = WaterYOLODetector(model_path="yolov8n.pt", device="0") # 若无GPU改为 "cpu"

# 模拟的 RTSP 摄像头流字典，实际中可动态增删
camera_streams = {
    "cam_001": "rtsp://admin:password@192.168.1.100:554/stream1",
    "cam_002": "rtsp://admin:password@192.168.1.101:554/stream1"
}

# Django 后端接收告警的 Webhook 地址
DJANGO_WEBHOOK_URL = "http://localhost:8000/api/ai-alerts/webhook/"

class AlertPayload(BaseModel):
    camera_id: str
    alert_type: str
    confidence: float
    timestamp: str
    image_base64: str

def send_alert_to_django(payload: dict):
    """后台任务：将告警推送到 Django"""
    try:
        # 实际生产中建议加入重试机制或使用 Celery/RabbitMQ
        response = requests.post(DJANGO_WEBHOOK_URL, json=payload, timeout=5)
        if response.status_code == 200:
            print(f"[Success] Alert sent to Django for camera {payload['camera_id']}")
        else:
            print(f"[Failed] Django responded with {response.status_code}")
    except Exception as e:
        print(f"[Error] Failed to send alert: {e}")

def process_camera_stream(camera_id: str, rtsp_url: str):
    """
    持续处理摄像头视频流 (运行在独立线程中)
    """
    print(f"[Info] Starting stream processing for {camera_id}: {rtsp_url}")
    cap = cv2.VideoCapture(rtsp_url)
    
    # 设置缓冲区大小，避免延迟累积
    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)

    frame_count = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            print(f"[Warning] Failed to read frame from {camera_id}. Reconnecting...")
            cap.release()
            time.sleep(3) # 断流重连等待
            cap = cv2.VideoCapture(rtsp_url)
            continue

        frame_count += 1
        # 性能优化：不需要每帧都检测，例如每 15 帧检测一次 (假设 30fps，即每秒检测 2 次)
        if frame_count % 15 != 0:
            continue

        # 1. YOLO 推理
        result = detector.process_frame(frame, conf_threshold=0.6)
        
        # 2. 判断是否触发告警
        if result["detections"]:
            for det in result["detections"]:
                # 构建告警数据
                import datetime
                payload = {
                    "camera_id": camera_id,
                    "alert_type": det["class"],
                    "confidence": det["confidence"],
                    "timestamp": datetime.datetime.now().isoformat(),
                    "image_base64": detector.encode_image_to_base64(result["annotated_frame"])
                }
                # 3. 异步发送给 Django
                # 注意：在真实高并发场景中，这里应推送到 Redis 队列，由专门的 worker 消费
                send_alert_to_django(payload)
                
        # (可选) 如果你想提供实时视频流 API 给前端，可以将 result["annotated_frame"] 编码为 MJPEG 
        # 但生产环境更推荐使用 WebRTC 或专门的流媒体服务器 (如 ZLMediaKit)

@app.on_event("startup")
async def startup_event():
    """服务启动时，自动开启所有摄像头的后台检测线程"""
    for cam_id, rtsp_url in camera_streams.items():
        thread = threading.Thread(target=process_camera_stream, args=(cam_id, rtsp_url), daemon=True)
        thread.start()
        print(f"[Info] Thread started for {cam_id}")

@app.get("/health")
def health_check():
    return {"status": "healthy", "model_loaded": True}

# 启动命令: uvicorn main:app --host 0.0.0.0 --port 8001

四、 与 Django 后端的联动 (Webhook 接收)

在 Django 端，你需要创建一个接口来接收 YOLO 服务推送的告警，并保存入库或触发短信通知。

Django views.py 示例:

python 复制代码 
from rest_framework.decorators import api_view, permission_classes
from rest_framework.permissions import AllowAny
from rest_framework.response import Response
from django.utils import timezone
from .models import AIAlertLog # 假设你有一个告警日志模型
import base64
import os

@api_view(['POST'])
@permission_classes([AllowAny]) # 实际生产中应使用内部 Token 验证
def ai_alert_webhook(request):
    try:
        data = request.data
        camera_id = data.get('camera_id')
        alert_type = data.get('alert_type')
        image_b64 = data.get('image_base64')
        
        # 1. 保存截图到本地或 OSS/MinIO
        if image_b64:
            img_data = base64.b64decode(image_b64)
            filename = f"alert_{camera_id}_{timezone.now().strftime('%Y%m%d_%H%M%S')}.jpg"
            filepath = os.path.join('media/ai_alerts', filename)
            os.makedirs(os.path.dirname(filepath), exist_ok=True)
            with open(filepath, 'wb') as f:
                f.write(img_data)
            
            # 更新数据中的图片路径，移除 base64 以节省数据库空间
            data['image_url'] = f"/media/ai_alerts/{filename}"
            del data['image_base64']

        # 2. 存入数据库
        alert = AIAlertLog.objects.create(
            camera_id=camera_id,
            alert_type=alert_type,
            confidence=data.get('confidence'),
            alert_time=timezone.now(),
            image_url=data.get('image_url'),
            is_handled=False
        )
        
        # 3. (可选) 触发 Celery 异步任务发送短信/微信通知
        # send_sms_alert.delay(alert.id)

        return Response({"status": "success", "alert_id": alert.id})
    except Exception as e:
        return Response({"status": "error", "message": str(e)}, status=400)

五、 针对水利场景的进阶优化建议

  1. 模型微调 (Fine-tuning) ：

    官方的 yolov8n.pt 只能识别通用物体（如 person, cup）。为了精准识别"水面漂浮物"、"未穿救生衣"、"特定型号的水位尺"，你需要：

    • 收集水利场景图片（可使用 Roboflow 平台标注）。
    • 使用 yolo task=detect mode=train data=your_dataset.yaml epochs=100 imgsz=640 训练专属模型。
    • 将生成的 best.pt 替换到 detector.py 中。

  2. 多路视频流性能瓶颈 ：

    上述代码使用 Python 多线程处理 RTSP。如果摄像头超过 4-8 路，Python 的 GIL 和 OpenCV 的解码会成为瓶颈。

    • 中级方案 ：使用 multiprocessing 为每个摄像头分配独立进程。
    • 高级方案 (生产推荐) ：放弃 Python 直接拉流，部署专业的流媒体网关（如 ZLMediaKit 或 SRS ）。YOLO 服务只通过 HTTP-FLV 或 共享内存 从网关获取抽帧图像，或者直接使用 NVIDIA DeepStream SDK 进行硬件级多路并发推理。

  3. 区域入侵检测 (ROI) ：

    水利场景中，并非画面里出现"人"就告警，而是"人进入了危险区域（如大坝边缘、闸门附近）"。

    • 实现 ：在 Django 后台配置每个摄像头的多边形坐标 (Polygon ROI)。YOLO 检测到 person 后，计算其 bbox 中心点是否落在该多边形内，只有落在内部才触发 Webhook 告警。

  4. 容器化部署 (Docker) ：

    为了方便在边缘计算盒子或服务器上部署，提供 Dockerfile：

    dockerfile 复制代码 
    FROM python:3.9-slim
# 安装系统级 OpenCV 依赖
RUN apt-get update && apt-get install -y libgl1-mesa-glx libglib2.0-0
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
# 预下载模型，避免容器启动时下载失败
RUN python -c "from ultralytics import YOLO; YOLO('yolov8n.pt')"
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8001"]
相关推荐
动物园猫1 小时前
无人机行人精准检测数据集分享（适用于YOLO系列深度学习分类检测任务）
深度学习·yolo·无人机
深度学习lover2 小时前
＜数据集＞yolo月球陨石坑识别＜目标检测＞
人工智能·yolo·目标检测·计算机视觉·数据集·月球陨石坑识别
稷下元歌4 小时前
aifei学习前置基础：全套完整教程：Anaconda 安装→环境配置→YOLOv8+OpenCV 安装 + OpenCV 实操 + 标注→训练→导出→部署
opencv·学习·yolo
深度学习lover14 小时前
＜数据集＞yolo樱桃识别＜目标检测＞
人工智能·深度学习·yolo·目标检测·计算机视觉·数据集·樱桃识别
王哈哈^_^1 天前
【源码教程+数据集】农作物分类检测数据集 10712 张，农作物分类检测系统实战教程
人工智能·算法·yolo·目标检测·计算机视觉·毕业设计·数据集
YOLO数据集集合1 天前
输电线缺陷目标检测｜无人机电力巡检深度学习数据集｜电网线缆散股智能识别数据
人工智能·深度学习·yolo·目标检测·无人机
王哈哈^_^1 天前
YOLO分类任务训练教程：从数据准备到模型部署全流程
人工智能·yolo·计算机视觉·分类·数据挖掘
动物园猫1 天前
无人机角度的道路损害检测数据集分享（适用于YOLO系列深度学习分类检测任务）
深度学习·yolo·无人机
victory_li1 天前
OpenVINO + Yolov26 + C# + .net framework4.8实现分类推理
yolo·c#·openvino
热门推荐
01GitHub 镜像站点022026 年 AI 编程工具终极横评：Cursor vs Claude Code vs Copilot vs Windsurf03【AI】2026 年具身智能模型和世界模型总结04Codex 下载安装指南：Windows 和 macOS 官方版下载05Codex 桌面端更新后 Chrome 插件和 Computer Use 不可用，怎么排查和修复06【踩坑记录 | 第一篇】微软商店无法使用时，如何手动安装 OpenAI Codex？附`.msix`文件系统错误解决方法07CC-Switch 下载、安装与使用配置指南【2026.5.29】08Codex 接入 DeepSeek API 完整配置文档09CC-Switch & Claude 基于 Linux 服务器安装使用指南10裂开！ChatGPT 居然开始要手机号验证，附详细解决方法