使用 FastAPI 上传图片并传递给 YOLOv8 模型

本文档旨在指导开发者如何使用 FastAPI 框架构建一个 REST API 接口，该接口能够接收图片上传，并将图片数据传递给 YOLOv8 模型进行处理。我们将重点介绍如何处理上传的图片文件，并将其转换为 YOLOv8 模型能够接受的格式，解决直接传递字节数据导致的 “Unsupported image type” 错误。

问题背景

在使用 FastAPI 构建 REST API 接口时，经常需要处理图片上传的场景。当将上传的图片直接以字节流的形式传递给 YOLOv8 模型时，可能会遇到 “Unsupported image type” 错误。这是因为 YOLOv8 模型通常需要特定格式的图像数据，例如 PIL (Pillow) 库中的 Image 对象，而不是原始的字节流。

解决方案

解决此问题的关键在于将上传的图片字节流转换为 YOLOv8 模型能够接受的图像格式。以下是一个使用 FastAPI 和 PIL 库实现此功能的示例代码：

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.get("/")
def home():
return {"health_check": "OK", "model_version": 0.01}
@app.post("/predict/")
async def upload_file(file: UploadFile):
try:
content_byte = await file.read() # 读取上传文件的内容为字节流
content_image = Image.open(io.BytesIO(content_byte)) # 将字节流转换为 PIL Image 对象
result = model.load_yolo_v8(content_image) # 调用 YOLOv8 模型进行预测
return {"result": result}
except Exception as e:
return {"error": str(e)}

代码解释：

1. 导入必要的库： 导入 FastAPI、File、UploadFile 用于处理文件上传，PIL (Pillow) 用于图像处理，io 用于处理字节流，以及你的 YOLOv8 模型所在的模块。
2. 定义 API 接口： 使用 @app.post(“/predict/”) 定义一个 POST 请求接口，用于接收图片上传。
3. 读取文件内容： 使用 content_byte = await file.read() 读取上传文件的内容，得到字节流数据。注意这里需要使用await关键字，因为file.read()是一个异步操作。
4. 转换为 PIL Image 对象： 使用 content_image = Image.open(io.BytesIO(content_byte)) 将字节流数据转换为 PIL Image 对象。io.BytesIO 用于将字节流模拟成一个文件对象，然后 Image.open 函数可以从这个文件对象中读取图像数据。
5. 调用 YOLOv8 模型： 使用 result = model.load_yolo_v8(content_image) 调用你的 YOLOv8 模型，并将 PIL Image 对象作为输入传递给模型。 请注意，这里假设你的model.load_yolo_v8函数接受PIL Image对象作为输入。如果你的模型需要其他格式，请相应地进行转换。
6. 返回结果： 将模型的预测结果以 JSON 格式返回。

注意事项：

• 确保已安装必要的库：pip install fastapi uvicorn pillow
• 确保 app.model.model 模块中包含正确的 YOLOv8 模型加载和预测逻辑。
• 根据你的 YOLOv8 模型的要求，可能需要对 PIL Image 对象进行进一步的处理，例如调整大小、转换为特定颜色空间等。
• 在生产环境中，建议对上传的文件进行安全检查，例如验证文件类型、大小等，以防止恶意攻击。
• 异常处理部分可以根据实际需求进行更详细的错误信息返回。

总结

通过使用 FastAPI 和 PIL 库，我们可以轻松地构建一个 REST API 接口，用于接收图片上传，并将图片数据转换为 YOLOv8 模型能够接受的格式。这种方法可以有效地解决 “Unsupported image type” 错误，并使你能够更好地利用 YOLOv8 模型进行图像处理任务。记得根据你的模型需求调整代码，并添加必要的安全措施。