本文档介绍了如何使用 FastAPI 构建一个 REST API 接口，该接口能够接收图像文件，并将其传递给 YOLOv8 模型进行预测。重点讲解如何处理上传的图像数据，将其转换为 YOLOv8 模型所支持的格式，并展示了完整的代码示例，帮助开发者快速搭建图像预测服务。
图像上传与处理
在使用 YOLOv8 模型构建图像预测 API 时，一个常见的挑战是如何正确地接收和处理上传的图像数据。默认情况下，FastAPI 会将上传的文件以字节流 (bytes) 的形式传递给后端。然而，YOLOv8 模型通常需要图像文件路径或图像对象作为输入，直接传递字节流会导致模型无法识别图像类型。
以下代码展示了如何使用 Python Imaging Library (PIL) 将上传的字节流转换为 YOLOv8 模型可以接受的图像对象：
from fastapi import FastAPI, File, UploadFile
from pydantic import BaseModel
from PIL import Image
import io
import app.model.model as model # 假设你的 YOLOv8 模型在 app.model.model 模块中
app = FastAPI()
class PredictionOut(BaseModel):
result: list
@app.post("/predict/")
async def upload_file(file: UploadFile):
try:
content_byte = await file.read() # 使用 await 读取文件内容
content_image = Image.open(io.BytesIO(content_byte))
result = model.load_yolo_v8(content_image) # 修改为模型接受图像对象的方法
return {"result": result}
except Exception as e:
return {"error": str(e)}代码解释：
- await file.read(): 使用 await 异步读取上传的文件内容，得到字节流数据。
- io.BytesIO(content_byte): 将字节流数据封装到 io.BytesIO 对象中，使其可以像文件一样被 PIL 处理。
- Image.open(io.BytesIO(content_byte)): 使用 PIL 的 Image.open() 方法从 io.BytesIO 对象中读取图像数据，并创建一个 Image 对象。
- model.load_yolo_v8(content_image): 将 PIL 的 Image 对象传递给你的 YOLOv8 模型进行预测。 请注意，你需要根据你的 model.load_yolo_v8 函数的实际参数修改此处的调用方式。 例如，如果你的模型需要图像路径，你可能需要先将 Image 对象保存到临时文件，然后将文件路径传递给模型。
注意事项
- 异步读取文件： 使用 await file.read() 可以异步读取文件内容，避免阻塞主线程，提高 API 的响应速度。
- 错误处理： 代码中包含了 try...except 块来捕获可能发生的异常，并返回错误信息。 在实际应用中，你应该根据具体情况进行更详细的错误处理。
- 模型输入： 请务必确认你的 YOLOv8 模型接受的输入类型。如果模型需要图像文件路径，你需要将 PIL 的 Image 对象保存到临时文件，并将文件路径传递给模型。
- 依赖安装： 确保你已经安装了必要的依赖包，例如 fastapi, uvicorn, Pillow (PIL)。 可以使用 pip install fastapi uvicorn Pillow 命令进行安装。
- 模型加载： 请根据你的模型加载方式修改 model.load_yolo_v8 函数的调用方式。
总结
通过使用 PIL 库，我们可以轻松地将 FastAPI 上传的图像字节流转换为 YOLOv8 模型可以接受的图像对象，从而实现图像预测功能。 在实际应用中，请根据你的模型和需求进行适当的调整和优化。 记住要处理异常，并确保你的模型输入类型正确。
