27、利用FastAPI和WebSocket实现实时目标检测系统

利用FastAPI和WebSocket实现实时目标检测系统

1. 实现基础目标检测系统

在构建实时目标检测系统时，首先要实现一个基础的目标检测系统。关键在于正确使用 UploadFile 和 File 依赖项来获取上传的文件。若需要复习相关内容，可参考相关资料。获取文件后，将其实例化为合适的Pillow图像对象，并调用预测模型。同时，别忘了在生命周期处理程序中加载模型。

可以使用以下Uvicorn命令运行示例：

(venv) $ uvicorn chapter13.chapter13_api:app

使用HTTPie上传之前用过的咖啡店图片进行测试：

$ http --form POST http://localhost:8000/object-detection image@./assets/coffee-shop.jpg

返回结果如下：

{
    "objects": [
        {
            "box": [659.8709716796875, 592.8882446289062, 792.0460815429688, 840.2132568359375],
            "label": "person"
        },
        {
            "box": [873.5499267578125, 875.7918090820312, 1649.1378173828125, 1296.362548828125],
            "label": "couch"
        }
    ]
}

从结果可知，成功获取了检测到的对象列表，每个对象都有其边界框和标签。至此，目标检测系统已可作为Web服务器使用，但我们的目标是构建实时系统，而WebSocket能帮助我们处理图像流。

2. 实现WebSocket进行图像流目标检测

2.1 WebSocket的优势

WebSocket的主要优势之一是它在客户端和服务器之间打开了一个全双工通信通道。连接建立后，消息可以快速传递，无需经过HTTP协议的所有步骤，因此更适合实时发送大量数据。

2.2 实现思路与挑战

要实现一个既能接受图像数据又能对其进行目标检测的WebSocket端点，主要挑战是处理“背压”现象。由于运行检测算法需要时间，浏览器发送的图像数量会超过服务器的处理能力。因此，需要使用一个有限大小的队列（或缓冲区），并在处理过程中丢弃一些图像，以实现接近实时的流处理。

2.3 代码实现

以下是具体的代码实现：

import asyncio
from fastapi import WebSocket
from PIL import Image
import io

# 假设object_detection是已经定义好的目标检测模型
async def receive(websocket: WebSocket, queue: asyncio.Queue):
    while True:
        bytes = await websocket.receive_bytes()
        try:
            queue.put_nowait(bytes)
        except asyncio.QueueFull:
            pass

async def detect(websocket: WebSocket, queue: asyncio.Queue):
    while True:
        bytes = await queue.get()
        image = Image.open(io.BytesIO(bytes))
        objects = object_detection.predict(image)
        await websocket.send_json(objects.dict())

from fastapi import FastAPI

app = FastAPI()

@app.websocket("/object-detection")
async def ws_object_detection(websocket: WebSocket):
    await websocket.accept()
    queue: asyncio.Queue = asyncio.Queue(maxsize=1)
    receive_task = asyncio.create_task(receive(websocket, queue))
    detect_task = asyncio.create_task(detect(websocket, queue))
    try:
        done, pending = await asyncio.wait(
            {receive_task, detect_task},
            return_when=asyncio.FIRST_COMPLETED,
        )
        for task in pending:
            task.cancel()
        for task in done:
            task.result()
    except WebSocketDisconnect:
        pass

这里定义了两个任务： receive 和 detect 。 receive 任务等待来自WebSocket的原始字节数据，并将其放入队列； detect 任务从队列中取出数据，进行目标检测，并将结果发送回客户端。

asyncio.Queue 对象是关键，它允许我们在内存中对数据进行排队，并采用先进先出（FIFO）策略检索数据。通过设置队列的最大元素数量，可以限制处理的图像数量。

2.4 静态文件服务

在完整的实现中，还定义了两个额外的功能：一个索引端点，用于返回 index.html 文件；一个 StaticFiles 应用，挂载在 /assets 路径下。这两个功能使FastAPI应用能够直接提供HTML和JavaScript代码，方便浏览器在同一服务器上查询这些文件。这表明FastAPI不仅适合构建REST API，还能很好地服务HTML和静态文件。

3. 从浏览器通过WebSocket发送图像流

3.1 实现概述

后端准备好后，接下来要从浏览器捕获图像并通过WebSocket发送。这主要涉及JavaScript代码，虽然超出了部分范围，但为了使应用完整运行是必要的。

3.2 具体步骤

3.2.1 启用摄像头输入并列出可用摄像头

window.addEventListener('DOMContentLoaded', (event) => {
    const video = document.getElementById('video');
    const canvas = document.getElementById('canvas');
    const cameraSelect = document.getElementById('camera-select');
    let socket;

    // 列出可用摄像头并填充选择框
    navigator.mediaDevices.getUserMedia({ audio: true, video: true }).then(() => {
        navigator.mediaDevices.enumerateDevices().then((devices) => {
            for (const device of devices) {
                if (device.kind === 'videoinput' && device.deviceId) {
                    const deviceOption = document.createElement('option');
                    deviceOption.value = device.deviceId;
                    deviceOption.innerText = device.label;
                    cameraSelect.appendChild(deviceOption);
                }
            }
        });
    });
});

3.2.2 启动目标检测

用户提交表单后，调用 startObjectDetection 函数：

// 提交表单时在所选摄像头上启动目标检测
document.getElementById('form-connect').addEventListener('submit', (event) => {
    event.preventDefault();
    // 关闭之前的套接字（如果有）
    if (socket) {
        socket.close();
    }
    const deviceId = cameraSelect.selectedOptions[0].value;
    socket = startObjectDetection(video, canvas, deviceId);
});

3.2.3 startObjectDetection 函数实现

const startObjectDetection = (video, canvas, deviceId) => {
    const socket = new WebSocket(`ws://${location.host}/object-detection`);
    let intervalId;

    // 连接打开
    socket.addEventListener('open', function () {
        // 开始从设备读取视频
        navigator.mediaDevices.getUserMedia({
            audio: false,
            video: {
                deviceId,
                width: { max: 640 },
                height: { max: 480 },
            },
        }).then(function (stream) {
            video.srcObject = stream;
            video.play().then(() => {
                // 调整覆盖画布大小以适应视频大小
                canvas.width = video.videoWidth;
                canvas.height = video.videoHeight;

                // 每42毫秒在WebSocket中发送一张图像
                intervalId = setInterval(() => {
                    // 创建一个虚拟画布来绘制当前视频图像
                    const canvas = document.createElement('canvas');
                    const ctx = canvas.getContext('2d');
                    canvas.width = video.videoWidth;
                    canvas.height = video.videoHeight;
                    ctx.drawImage(video, 0, 0);
                    // 将其转换为JPEG并发送到WebSocket
                    canvas.toBlob((blob) => socket.send(blob), 'image/jpeg');
                }, IMAGE_INTERVAL_MS);
            });
        });
    });

    // 监听消息
    socket.addEventListener('message', function (event) {
        drawObjects(video, canvas, JSON.parse(event.data));
    });

    // 关闭时停止间隔和视频读取
    socket.addEventListener('close', function () {
        window.clearInterval(intervalId);
        video.pause();
    });

    return socket;
};

这里限制了视频输入的大小为640x480像素，以避免过大的图像对服务器造成压力。同时，设置每42毫秒发送一张图像，大约相当于每秒24帧。

3.3 流程图

graph TD;
    A[页面加载] --> B[列出可用摄像头];
    B --> C[用户选择摄像头并提交表单];
    C --> D[启动WebSocket连接];
    D --> E[开始读取视频];
    E --> F[定时捕获图像并发送];
    F --> G[接收服务器返回结果];
    G --> H[绘制检测结果];

4. 在浏览器中显示目标检测结果

4.1 实现思路

要在浏览器中显示目标检测结果，可使用 <canvas> 元素。通过CSS将其覆盖在视频上，然后根据服务器发送的矩形坐标绘制绿色矩形和标签。

4.2 HTML代码

<body>
    <div class="container">
        <h1 class="my-3">Chapter 13 - Real time object detection</h1>
        <form id="form-connect">
            <div class="input-group mb-3">
                <select id="camera-select"></select>
                <button class="btn btn-success" type="submit" id="button-start">Start</button>
            </div>
        </form>
        <div class="position-relative" style="width: 640px; height: 480px;">
            <video id="video"></video>
            <canvas id="canvas" class="position-absolute top-0 start-0"></canvas>
        </div>
    </div>
    <script src="/assets/script.js"></script>
</body>

4.3 JavaScript代码

const drawObjects = (video, canvas, objects) => {
    const ctx = canvas.getContext('2d');
    ctx.width = video.videoWidth;
    ctx.height = video.videoHeight;

    ctx.beginPath();
    ctx.clearRect(0, 0, ctx.width, ctx.height);
    for (const object of objects.objects) {
        const [x1, y1, x2, y2] = object.box;
        const label = object.label;
        ctx.strokeStyle = '#49fb35';
        ctx.beginPath();
        ctx.rect(x1, y1, x2 - x1, y2 - y1);
        ctx.stroke();
        ctx.font = 'bold 16px sans-serif';
        ctx.fillStyle = '#ff0000';
        ctx.fillText(label, x1 - 5 , y1 - 5);
    }
};

在 drawObjects 函数中，首先清除之前的绘制结果，然后遍历所有检测到的对象，绘制矩形和标签。

4.4 运行系统

使用以下命令启动系统：

(venv) $ uvicorn chapter13.websocket_object_detection.app:app

在浏览器中访问 http://localhost:8000 ，选择摄像头并点击“Start”，即可看到视频输出和目标检测结果。

通过以上步骤，我们成功实现了一个基于FastAPI和WebSocket的实时目标检测系统，将Python系统的智能带到了用户的浏览器中，为用户提供了接近实时的交互体验。

5. 实时目标检测系统总结

5.1 系统成果

通过一系列的操作，我们成功构建了一个基于FastAPI和WebSocket的实时目标检测系统。利用Hugging Face社区提供的预训练模型，快速实现了目标检测功能，并将其集成到WebSocket端点中。借助现代JavaScript API，实现了从浏览器发送视频输入并直接显示算法结果，为用户带来了更加互动的体验。

5.2 工具优势

FastAPI等强大工具在这个项目中发挥了重要作用，使我们能够在短时间内获得成果，并且代码易于理解。例如，FastAPI不仅适合构建REST API，还能很好地服务HTML和静态文件，为系统的完整性提供了保障。

5.3 潜在问题

在之前的示例和项目中，我们假设所使用的机器学习模型运行速度足够快，可以直接在API端点或WebSocket任务中运行。但实际情况并非总是如此，某些复杂算法可能需要几分钟才能完成。如果直接在API端点中运行这类算法，用户需要长时间等待响应，这不仅体验不佳，还会阻塞服务器，影响其他用户使用API。

6. 分布式文本到图像AI系统的引入

6.1 传统API处理的局限性

在之前构建的API中，所有操作都在请求处理过程中完成，用户需要等待服务器完成所有定义的操作才能获得响应。然而，这种方式并非总是理想或可行的。

例如，在Web应用中发送电子邮件通知，服务器需要向邮件服务器发送请求，这个操作可能需要几毫秒。如果在请求处理过程中进行此操作，响应会延迟，而用户实际上并不关心邮件的发送方式和时间。

另一个例子是用户请求复杂的数据导出或使用重型AI模型时，这些操作无法在合理时间内完成。如果在请求处理程序中执行这些操作，会阻塞服务器进程，导致服务器响应缓慢。而且，一些网络基础设施（如代理或浏览器）有严格的超时设置，长时间无响应的操作可能会被取消。

6.2 解决方案：Web - 队列 - 工作者架构

为了解决上述问题，我们引入了Web - 队列 - 工作者架构。这种架构将最昂贵、耗时的操作推迟到后台进程（工作者）中进行。

以下是该架构的主要组成部分：
| 组成部分 | 描述 |
| ---- | ---- |
| Web | 负责接收用户请求，处理简单的验证和路由，将复杂任务放入队列 |
| 队列 | 作为任务的缓冲区，存储待处理的任务，确保任务按顺序执行 |
| 工作者 | 从队列中取出任务，执行复杂的操作，如AI模型推理、数据处理等 |

6.3 示例项目：基于Stable Diffusion模型的文本到图像AI系统

为了展示这种架构的实际应用，我们将构建一个基于Stable Diffusion模型的文本到图像AI系统。用户输入文本提示，系统将根据提示生成相应的图像。

7. 分布式文本到图像AI系统的工作流程

7.1 流程图

graph LR;
    A[用户输入文本提示] --> B[Web服务器接收请求];
    B --> C[将任务放入队列];
    C --> D[工作者从队列取出任务];
    D --> E[工作者使用Stable Diffusion模型生成图像];
    E --> F[工作者将结果返回给Web服务器];
    F --> G[Web服务器将图像返回给用户];

7.2 详细步骤

7.2.1 用户请求

用户在前端界面输入文本提示，点击生成按钮，请求发送到Web服务器。

7.2.2 Web服务器处理

Web服务器接收到请求后，对请求进行简单的验证，确保输入的合法性。然后将任务信息（如文本提示）放入队列中，并向用户返回一个任务ID，表示任务已接收。

7.2.3 队列管理

队列负责存储待处理的任务，采用先进先出（FIFO）的策略。工作者会定期从队列中取出任务进行处理。

7.2.4 工作者处理

工作者从队列中取出任务，加载Stable Diffusion模型，根据文本提示生成图像。这个过程可能需要一些时间，取决于模型的复杂度和计算资源。

7.2.5 结果返回

工作者完成图像生成后，将结果（图像数据）返回给Web服务器。Web服务器将结果存储在合适的位置，并根据任务ID将图像返回给用户。

8. 分布式系统的优势

8.1 提高响应速度

通过将复杂任务交给工作者处理，Web服务器可以快速响应用户请求，避免长时间的等待，提高用户体验。

8.2 增强系统可扩展性

可以根据实际需求增加或减少工作者的数量，以应对不同的负载。当用户请求量增加时，增加工作者数量可以提高系统的处理能力。

8.3 资源有效利用

工作者可以在独立的计算资源上运行，充分利用服务器的计算能力，避免资源浪费。

综上所述，引入分布式文本到图像AI系统可以解决传统API处理复杂任务时的局限性，为用户提供更加高效、稳定的服务。通过合理的架构设计和资源分配，可以实现系统的高性能和可扩展性。