26、利用FastAPI构建高效预测API与实时目标检测系统

利用FastAPI构建高效预测API与实时目标检测系统

1. 使用Joblib缓存预测结果

在开发预测API时，为了提高效率，我们可以使用Joblib来缓存预测结果。在预测方法中，我们调用外部的预测函数，并将模型和输入文本作为参数传递。之后，我们只需检索相应的类别名称并构建一个 PredictionOutput 对象。

以下是相关代码示例：

category = self.targets[prediction]
return PredictionOutput(category=category)

代码链接： chapter12_caching.py

我们还添加了一个 delete/cache 路由，以便通过HTTP请求清除整个Joblib缓存。示例代码如下：

@app.post("/prediction")
def prediction(
    output: PredictionOutput = Depends(newgroups_model.predict),
) -> PredictionOutput:
    return output

@app.delete("/cache", status_code=status.HTTP_204_NO_CONTENT)
def delete_cache():
    memory.clear()

代码链接： chapter12_caching.py

memory 对象的 clear 方法会删除磁盘上所有的Joblib缓存文件。这样，当我们使用相同的输入进行两次请求时，第二次响应将显示缓存的结果。虽然在这个示例中，模型运行速度很快，执行时间的差异不太明显，但对于更复杂的模型，缓存可以显著提高性能。

2. 选择标准函数还是异步函数

FastAPI支持异步I/O，这对应用程序的性能有很大的提升。然而，在某些情况下，比如使用Joblib时，它是同步实现的，并且会执行长时间的I/O操作（如读写硬盘上的缓存文件），这会阻塞进程，导致在操作进行时无法响应其他请求。

为了解决这个问题，FastAPI实现了一个巧妙的机制：如果将路径操作函数或依赖项定义为标准的非异步函数，它将在单独的线程中运行。这样，像同步文件读取这样的阻塞操作就不会阻塞主进程，类似于异步操作。

为了更好地理解，我们进行一个简单的实验。构建一个包含三个端点的虚拟FastAPI应用程序：
- /fast ：直接返回响应。
- /slow-async ：定义为异步的路径操作，创建一个需要运行10秒的同步阻塞操作。
- /slow-sync ：定义为标准方法的路径操作，创建一个需要运行10秒的同步阻塞操作。

以下是代码示例：

import time
from fastapi import FastAPI

app = FastAPI()

@app.get("/fast")
async def fast():
    return {"endpoint": "fast"}

@app.get("/slow-async")
async def slow_async():
    """Runs in the main process"""
    time.sleep(10)  # Blocking sync operation
    return {"endpoint": "slow-async"}

@app.get("/slow-sync")
def slow_sync():
    """Runs in a thread"""
    time.sleep(10)  # Blocking sync operation
    return {"endpoint": "slow-sync"}

代码链接： chapter12_async_not_async.py

我们可以使用Uvicorn运行这个应用程序：

(venv) $ uvicorn chapter12.chapter12_async_not_async:app

然后，我们进行以下实验：
1. 在第一个终端中，向 /slow-async 端点发送请求：

$ http GET http://localhost:8000/slow-async

2. 不等待响应，在第二个终端中，向 /fast 端点发送请求：

$ http GET http://localhost:8000/fast

你会发现，在收到 /fast 端点的响应之前，你需要等待10秒。这意味着 /slow-async 阻塞了进程，阻止了服务器在操作进行时响应其他请求。

接下来，我们对 /slow-sync 端点进行同样的实验：

$ http GET http://localhost:8000/slow-sync
$ http GET http://localhost:8000/fast

这次，你会立即得到 /fast 端点的响应，而无需等待 /slow-sync 完成。因为它被定义为标准的非异步函数，FastAPI会在单独的线程中运行它，从而避免阻塞。

在开发FastAPI应用程序时，选择标准函数还是异步函数的经验法则如下：
| 情况 | 选择建议 |
| ---- | ---- |
| 函数不涉及长时间的I/O操作（如文件读取、网络请求等） | 定义为异步函数 |
| 函数涉及I/O操作 | 1. 尝试选择支持异步I/O的库，函数定义为异步函数
2. 如果无法使用异步库（如Joblib缓存），定义为标准函数，FastAPI会在单独的线程中运行 |

由于Joblib在进行I/O操作时是完全同步的，我们将路径操作和依赖方法改为同步的标准方法。虽然在这个示例中，差异不太明显，但在处理较慢的操作（如将文件上传到云存储）时，这一点需要牢记。

3. 利用Hugging Face实现计算机视觉模型

在某些场景下，我们需要对输入流进行连续的预测，例如实时目标检测系统。FastAPI除了支持HTTP端点外，还支持WebSocket端点，这使得我们可以发送和接收数据流。我们可以利用Hugging Face的工具和预训练模型来构建一个实时目标检测系统。

3.1 安装所需库

首先，我们需要安装项目所需的所有库：

(venv) $ pip install "transformers[torch]" Pillow

transformers 库允许我们下载和运行预训练的机器学习模型，我们选择安装带有 torch 依赖项的版本。 Pillow 是一个广泛使用的Python图像处理库。

3.2 实现简单的目标检测脚本

在使用FastAPI之前，我们可以实现一个简单的脚本来运行目标检测算法，主要包括以下四个步骤：
1. 使用 Pillow 从磁盘加载图像。
2. 加载预训练的目标检测模型。
3. 在图像上运行模型。
4. 通过在检测到的对象周围绘制矩形来显示结果。

以下是具体的代码实现：

from pathlib import Path
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import YolosForObjectDetection, YolosImageProcessor

root_directory = Path(__file__).parent.parent
picture_path = root_directory / "assets" / "coffee-shop.jpg"
image = Image.open(picture_path)

image_processor = YolosImageProcessor.from_pretrained("hustvl/yolos-tiny")
model = YolosForObjectDetection.from_pretrained("hustvl/yolos-tiny")

inputs = image_processor(images=image, return_tensors="pt")
outputs = model(**inputs)

target_sizes = torch.tensor([image.size[::-1]])
results = image_processor.post_process_object_detection(
    outputs, target_sizes=target_sizes
)[0]

draw = ImageDraw.Draw(image)
font_path = root_directory / "assets" / "OpenSans-ExtraBold.ttf"
font = ImageFont.truetype(str(font_path), 24)
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
    if score > 0.7:
        box_values = box.tolist()
        label = model.config.id2label[label.item()]
        draw.rectangle(box_values, outline="red", width=5)
        draw.text(box_values[0:2], label, fill="red", font=font)
image.show()

代码链接： chapter13_object_detection.py

运行这个脚本：

(venv) $ python chapter13/chapter13_object_detection.py

第一次运行时，你会看到模型正在下载。预测可能需要几秒钟的时间，具体取决于你的计算机性能。完成后，结果图像将自动打开。

4. 实现用于单张图像目标检测的REST端点

在使用WebSocket之前，我们可以先实现一个经典的HTTP端点，用于接受图像上传并对其进行目标检测。与之前的示例相比，主要的区别在于图像的获取方式：我们从文件上传中获取图像，并将其转换为 Pillow 图像对象。

4.1 定义Pydantic模型

首先，我们定义Pydantic模型来正确构建预测模型的输出：

class Object(BaseModel):
    box: tuple[float, float, float, float]
    label: str

class Objects(BaseModel):
    objects: list[Object]

代码链接： chapter13_api.py

Object 模型表示单个检测到的对象，包含边界框的坐标和对象的标签。 Objects 模型是一个包含多个 Object 的列表。

4.2 实现FastAPI端点

以下是FastAPI端点的实现代码：

object_detection = ObjectDetection()

@contextlib.asynccontextmanager
async def lifespan(app: FastAPI):
    object_detection.load_model()
    yield

app = FastAPI(lifespan=lifespan)

@app.post("/object-detection", response_model=Objects)
async def post_object_detection(image: UploadFile = File(...)) -> Objects:
    image_object = Image.open(image.file)
    return object_detection.predict(image_object)

代码链接： chapter13_api.py

通过以上步骤，我们可以利用FastAPI构建一个高效的预测API和实时目标检测系统，为用户提供强大的服务。在实际应用中，我们可以根据具体需求对代码进行进一步的优化和扩展。

利用FastAPI构建高效预测API与实时目标检测系统

5. 通过WebSocket实现实时目标检测系统

在前面的内容中，我们已经完成了单张图像目标检测的REST端点实现。接下来，我们将利用FastAPI的WebSocket功能构建一个实时目标检测系统。在这个系统中，浏览器会将来自网络摄像头的图像流发送到WebSocket，我们的应用程序将运行目标检测算法，并将检测到的每个对象的坐标和标签发送回浏览器。

5.1 系统架构概述

整个实时目标检测系统的架构可以用以下mermaid流程图表示：

graph LR
    A[浏览器] -->|发送图像流| B[WebSocket]
    B -->|接收图像流| C[目标检测算法]
    C -->|运行检测| D[获取检测结果]
    D -->|返回结果| B
    B -->|发送结果| A

在这个流程中，浏览器负责采集图像并发送到WebSocket，服务器端的目标检测算法对图像进行处理，最后将结果返回给浏览器进行显示。

5.2 服务器端实现

在服务器端，我们需要实现一个WebSocket端点来处理图像流和返回检测结果。以下是一个简化的示例代码：

from fastapi import FastAPI, WebSocket
from PIL import Image
import io
from your_object_detection_module import detect_objects  # 假设这是我们的目标检测函数

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        try:
            data = await websocket.receive_bytes()
            image = Image.open(io.BytesIO(data))
            results = detect_objects(image)
            await websocket.send_json(results)
        except Exception as e:
            print(f"Error: {e}")
            break

在这个代码中，我们定义了一个WebSocket端点 /ws 。当客户端连接到这个端点时，我们接受连接并进入一个无限循环，不断接收客户端发送的图像数据。接收到图像数据后，我们将其转换为 Pillow 图像对象，并调用目标检测函数 detect_objects 进行检测。最后，将检测结果以JSON格式发送回客户端。

5.3 客户端实现

在客户端，我们可以使用JavaScript来实现图像采集和WebSocket通信。以下是一个简单的示例代码：

<!DOCTYPE html>
<html lang="en">

<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Real-time Object Detection</title>
</head>

<body>
    <video id="video" width="640" height="480" autoplay></video>
    <canvas id="canvas" width="640" height="480"></canvas>

    <script>
        const video = document.getElementById('video');
        const canvas = document.getElementById('canvas');
        const ctx = canvas.getContext('2d');

        navigator.mediaDevices.getUserMedia({ video: true })
          .then(stream => {
                video.srcObject = stream;
            })
          .catch(err => {
                console.error('Error accessing camera:', err);
            });

        const socket = new WebSocket('ws://localhost:8000/ws');

        socket.onopen = () => {
            setInterval(() => {
                ctx.drawImage(video, 0, 0, canvas.width, canvas.height);
                canvas.toBlob(blob => {
                    const reader = new FileReader();
                    reader.onloadend = () => {
                        const arrayBuffer = reader.result;
                        const uint8Array = new Uint8Array(arrayBuffer);
                        socket.send(uint8Array);
                    };
                    reader.readAsArrayBuffer(blob);
                }, 'image/jpeg');
            }, 100);
        };

        socket.onmessage = event => {
            const results = JSON.parse(event.data);
            // 这里可以根据检测结果在canvas上绘制矩形和标签
            console.log(results);
        };

        socket.onclose = () => {
            console.log('WebSocket connection closed');
        };
    </script>
</body>

</html>

在这个客户端代码中，我们首先使用 navigator.mediaDevices.getUserMedia 方法获取摄像头的视频流，并将其显示在 <video> 元素中。然后，我们创建一个WebSocket连接到服务器端的 /ws 端点。在连接建立后，我们每隔100毫秒将当前视频帧绘制到 <canvas> 上，并将其转换为二进制数据发送到服务器。当接收到服务器返回的检测结果时，我们可以根据结果在 <canvas> 上绘制矩形和标签。

6. 总结与展望

通过前面的内容，我们学习了如何利用Joblib缓存预测结果来提高API的效率，如何根据不同的情况选择标准函数或异步函数，以及如何使用Hugging Face的工具和预训练模型实现目标检测。同时，我们还实现了用于单张图像目标检测的REST端点和基于WebSocket的实时目标检测系统。

在实际应用中，我们可以根据具体需求对这些系统进行进一步的优化和扩展。例如，对于目标检测算法，我们可以选择更强大的预训练模型，或者对模型进行微调以提高检测的准确性。对于实时目标检测系统，我们可以优化图像采集和传输的效率，减少延迟。

以下是一个总结表格，概括了我们在本文中学习的主要内容：
| 内容 | 关键点 |
| ---- | ---- |
| 使用Joblib缓存预测结果 | 调用外部预测函数，添加 delete/cache 路由清除缓存 |
| 选择标准函数还是异步函数 | 根据是否涉及长时间I/O操作选择合适的函数类型 |
| 利用Hugging Face实现计算机视觉模型 | 安装 transformers 和 Pillow 库，实现目标检测脚本 |
| 实现用于单张图像目标检测的REST端点 | 定义Pydantic模型，实现FastAPI端点 |
| 实现实时目标检测系统 | 利用WebSocket处理图像流和返回检测结果 |

通过不断学习和实践，我们可以利用FastAPI构建出更加高效、强大的数据科学应用程序，为用户提供更好的服务体验。希望本文能为你在构建数据科学应用方面提供一些有用的参考和启发。