1. 使用Flask部署PyTorch图像分类模型

任务目标

构建一个Flask Web服务，加载PyTorch预训练的ResNet50模型，接收客户端上传的图像，返回Top-5分类结果（类别名+置信度）。

环境依赖

pip install flask torch torchvision pillow

完整实现代码

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.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
model.eval()

# 加载ImageNet类别标签
with open("imagenet_classes.json") as f:
    idx2label = json.load(f)

# 图像预处理：与训练时保持一致
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_bytes):
    """输入图像字节流，返回Top-5预测列表"""
    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    tensor = preprocess(img).unsqueeze(0)

    with torch.no_grad():
        outputs = model(tensor)
        probs = torch.nn.functional.softmax(outputs[0], dim=0)
        top5_probs, top5_indices = torch.topk(probs, 5)

    results = []
    for prob, idx in zip(top5_probs.tolist(), top5_indices.tolist()):
        results.append({
            "class": idx2label[str(idx)],
            "probability": round(prob, 4)
        })
    return results

@app.route("/predict", methods=["POST"])
def predict_api():
    if "image" not in request.files:
        return jsonify({"error": "请求中缺少image文件字段"}), 400

    file = request.files["image"]
    if file.filename == "":
        return jsonify({"error": "未选择文件"}), 400

    try:
        results = predict(file.read())
        return jsonify({"top5": results})
    except Exception as e:
        return jsonify({"error": f"图像处理失败: {str(e)}"}), 500

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

类别标签文件准备

下载ImageNet类别映射文件（约30KB），保存为imagenet_classes.json，与上述Python文件放在同一目录：

wget -O imagenet_classes.json https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt
# 将txt格式转换为JSON格式（执行以下Python命令）
python -c "import json; labels=open('imagenet_classes.txt').read().splitlines(); json.dump({str(i):labels[i] for i in range(len(labels))}, open('imagenet_classes.json','w'))"

运行与调用

1. 启动服务：python app.py，看到* Running on http://0.0.0.0:5000即启动成功。
2. 使用curl发送测试请求（替换为你的本地图片路径）：

curl -X POST -F "image=@/path/to/cat.jpg" http://localhost:5000/predict

验证方法与通过标准

• 正常响应：返回HTTP 200，JSON中包含top5数组，每个元素含class（字符串）和probability（0~1浮点数）。
• Top-1正确性：使用一张清晰的猫/狗图片测试，Top-1类别应为对应动物（如"tabby cat"、"golden retriever"），置信度通常大于0.5。
• 概率合法性：5个概率值均在0~1之间，且按降序排列。
• 错误处理：不传文件时返回400及错误信息；传非图片文件时返回500及错误信息。

常见问题

• 首次运行下载权重：ResNet50权重约100MB，首次执行时会自动下载到~/.cache/torch，需保持网络畅通。
• 内存不足：若GPU显存或CPU内存不足，可在predict()中增加torch.set_num_threads(1)限制CPU线程数。
• 类别索引错位：确保imagenet_classes.json的索引从0开始且与PyTorch模型输出顺序一致（均为ImageNet-1K标准顺序）。