2. 使用Flask创建模型API

任务目标

使用Flask框架封装一个已训练的图像分类模型（以PyTorch加载的ResNet18为例），创建REST API端点，接收客户端上传的图像文件，返回模型预测的类别名称及置信度。

环境依赖

需要安装：flask、torch、torchvision、pillow。若使用GPU推理则需对应CUDA版本PyTorch，CPU版本亦可运行。

完整实现代码

import io
import json
from flask import Flask, request, jsonify
from PIL import Image
import torch
import torchvision.transforms as transforms
from torchvision import models

app = Flask(__name__)

# 加载预训练模型并设为评估模式
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
model.eval()

# ImageNet类别标签（仅加载所需映射）
LABELS_URL = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
import urllib.request
with urllib.request.urlopen(LABELS_URL) as f:
    classes = [line.strip() for line in f.readlines()]

# 图像预处理：与训练时一致
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

def predict_image(image_bytes):
    """输入图像字节流，返回(类别名, 置信度)"""
    img = Image.open(io.BytesIO(image_bytes)).convert('RGB')
    input_tensor = preprocess(img).unsqueeze(0)  # 增加batch维度
    with torch.no_grad():
        output = model(input_tensor)
        probabilities = torch.nn.functional.softmax(output[0], dim=0)
        conf, pred_idx = torch.max(probabilities, dim=0)
    return classes[pred_idx.item()], conf.item()

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'error': '请求中缺少file字段'}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': '未选择文件'}), 400
    try:
        image_bytes = file.read()
        label, confidence = predict_image(image_bytes)
        return jsonify({
            'predicted_class': label,
            'confidence': round(confidence, 4)
        })
    except Exception as e:
        return jsonify({'error': f'图像处理失败: {str(e)}'}), 400

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

运行方式

1. 将上述代码保存为 app.py。
2. 在终端执行 python app.py，服务默认监听5000端口。
3. 首次运行会自动下载ResNet18权重（约45MB）和类别标签文件，需保持网络畅通。

调用与验证

使用curl命令发送测试请求（以本地任意一张dog.jpg为例）：

curl -X POST -F "file=@dog.jpg" http://127.0.0.1:5000/predict

预期输出（类别名称取决于测试图片内容）：

{"predicted_class":"golden retriever","confidence":0.9123}

通过条件：

• 返回HTTP 200且JSON中包含predicted_class与confidence字段；
• 类别名称与图片内容相符（如狗图片返回犬种名）；
• 置信度数值在0~1之间。

若上传非图像文件或缺少file字段，应返回HTTP 400及错误信息。

常见问题

• 端口被占用：修改app.run中的port参数。
• 内存不足：将model移到predict_image外部（已在代码中实现），避免每次请求重复加载。
• 首次请求慢：权重文件在服务启动时加载，首次请求仅需图像预处理与推理，通常秒级内完成。