使用 FastAPI 上传图片并应用于 YOLOv8 模型

第一段引用上面的摘要：

本文档旨在指导开发者如何使用 FastAPI 框架构建一个 REST API 接口，该接口能够接收上传的图片，并将其传递给 YOLOv8 模型进行处理。我们将详细介绍如何读取上传的图片文件，将其转换为 YOLOv8 模型可以接受的格式，并返回预测结果。通过本文的学习，你将掌握使用 FastAPI 处理图片上传，并将其应用于深度学习模型的关键技术。

构建 FastAPI 图片上传接口

FastAPI 是一个现代、高性能的 Python Web 框架，非常适合构建 API 接口。以下代码展示了如何创建一个接收图片上传的 API 接口：

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()  # 使用 await 读取文件内容
content_image = Image.open(io.BytesIO(content_byte))
result = model.predict_result(content_image)  # 假设你的模型需要 PIL Image 对象
return {"result": result}
except Exception as e:
return {"error": str(e)}

代码解释:

1. 导入必要的库:
   • FastAPI, File, UploadFile 来自 fastapi 用于构建 API 接口和处理文件上传。
   • BaseModel 来自 pydantic 用于定义数据模型。
   • Image 来自 PIL (Pillow) 用于图像处理。
   • io 用于处理内存中的字节流。
2. 创建 FastAPI 实例:
   • app = FastAPI() 创建一个 FastAPI 应用实例。
3. 定义预测结果模型:
   • PredictionOut(BaseModel) 定义 API 返回的数据结构。
4. 定义根路由:
   • @app.get(“/”) 定义一个 GET 请求的根路由，用于健康检查。
5. 定义图片上传路由:
   • @app.post(“/predict/”) 定义一个 POST 请求的 /predict/ 路由，用于接收图片上传。
   • async def upload_file(file: UploadFile) 定义处理上传文件的异步函数。
   • content_byte = await file.read() 使用 await 异步读取上传文件的内容，得到字节数据。
   • content_image = Image.open(io.BytesIO(content_byte)) 使用 PIL 库将字节数据转换为 Image 对象。
   • result = model.predict_result(content_image) 调用 YOLOv8 模型进行预测，假设 model.predict_result 函数接收 Image 对象作为输入。
   • return {“result”: result} 返回预测结果。
   • except Exception as e: return {“error”: str(e)} 捕获异常并返回错误信息。

关键步骤详解

• 读取文件内容: content_byte = await file.read() 这行代码使用 await 关键字异步读取上传文件的内容，返回的是字节数据。
• 转换为 PIL Image 对象: content_image = Image.open(io.BytesIO(content_byte)) 由于 YOLOv8 模型可能需要 PIL Image 对象作为输入，所以需要将字节数据转换为 PIL Image 对象。io.BytesIO(content_byte) 创建一个内存中的字节流，Image.open() 函数可以从这个字节流中读取图像数据。
• 调用 YOLOv8 模型: result = model.predict_result(content_image) 这行代码调用你的 YOLOv8 模型进行预测。请确保 model.predict_result 函数能够正确处理 PIL Image 对象，并返回预测结果。

注意事项

• 异步处理: 使用 async 和 await 关键字可以使 API 接口异步处理请求，提高并发性能。
• 错误处理: 在 try…except 块中捕获异常，并返回错误信息，可以提高 API 的健壮性。
• 模型输入格式: 请确保你的 YOLOv8 模型能够接受 PIL Image 对象作为输入。如果模型需要其他格式的输入，你需要进行相应的转换。
• 安装依赖: 确保安装了必要的 Python 库，如 fastapi, uvicorn, Pillow。 可以使用 pip install fastapi uvicorn Pillow 命令安装。
• 模型加载: 确保你的 app.model.model 模块正确加载了 YOLOv8 模型。

总结

通过本文，你学习了如何使用 FastAPI 构建一个图片上传接口，并将上传的图片转换为 YOLOv8 模型可以接受的格式。请务必根据你的实际情况修改代码，并进行充分的测试，以确保 API 接口能够正常工作。 记住，理解每个步骤背后的原理，才能更好地解决实际问题。