fastapi 封装yolo

### 封装YOLO模型于FastAPI中

#### 创建项目结构
为了有效地在FastAPI框架内封装YOLO模型并创建API服务，合理的项目结构至关重要。通常情况下，建议按照如下方式组织文件：

project/
│── app.py               # 主应用程序入口
├── models/              # 存放模型及其配置
│   └── yolo_model.pt    # 预训练好的YOLO模型权重文件
└── requirements.txt     # 依赖包列表

#### 安装必要的库
确保安装了所有必需的Python库来支持FastAPI和YOLO操作。这可以通过`requirements.txt`管理，并使用pip命令安装。

fastapi==0.95.0
uvicorn==0.20.0
torch>=1.8.0            # PyTorch版本需满足YOLOv5的要求
opencv-python-headless   # OpenCV用于图像处理而不启动GUI窗口
pydantic                 # 数据验证与解析

#### 加载YOLO模型
在`app.py`中定义一个函数用来加载预训练的YOLO模型实例。这里假设已经有一个名为`yolov5s.pt`的小规模YOLOv5模型作为例子[^1]。

import torch
from fastapi import FastAPI, File, UploadFile
from PIL import Image
import io

# 初始化FastAPI应用
app = FastAPI()

# 加载YOLO模型
model_path = "models/yolov5s.pt"
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path).to(device)

@app.post("/predict/")
async def predict(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))
    
    results = model(image)
    predictions = results.pandas().xyxy[0].to_dict(orient="records")

    return {"predictions": predictions}

上述代码片段展示了如何设置一个POST接口接受上传图片并通过已加载的YOLO模型执行预测任务。返回的结果是以JSON格式呈现的目标检测框位置和其他相关信息。

#### 运行服务器
完成以上步骤之后，就可以通过Uvicorn或其他ASGI服务器运行这个简单的Web服务了。可以在终端输入以下命令启动服务：

uvicorn app:app --reload

这样就成功地在一个轻量级且高效的异步环境中部署了一个基于YOLO的对象识别RESTful API。