基于 Flask 的图像分类预测案例解读

豆芽8192025-03-31 14:43

1

python 复制代码 
import io
import json
import flask
import torch
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
#from torchvision import transforms as T
from torchvision import transforms, models, datasets
from torch.autograd import Variable

# 初始化Flask app
app = flask.Flask(__name__)
model = None
use_gpu = False

# 加载模型进来
def load_model():
    """Load the pre-trained model, you can use your model just as easily.
    """
    global model
    #这里我们直接加载官方工具包里提供的训练好的模型（代码会自动下载）括号内参数为是否下载模型对应的配置信息
    model = models.resnet18()
    num_ftrs = model.fc.in_features
    model.fc = nn.Sequential(nn.Linear(num_ftrs, 102))  # 类别数自己根据自己任务来

    #print(model)
    checkpoint = torch.load('best.pth')
    model.load_state_dict(checkpoint['state_dict'])
    #将模型指定为测试格式
    model.eval()
    #是否使用gpu

    if use_gpu:
        model.cuda()

# 数据预处理
def prepare_image(image, target_size):
    """Do image preprocessing before prediction on any data.

    :param image:       original image
    :param target_size: target image size
    :return:
                        preprocessed image
    """
    #针对不同模型，image的格式不同，但需要统一至RGB格式
    if image.mode != 'RGB':
        image = image.convert("RGB")

    # Resize the input image and preprocess it.(按照所使用的模型将输入图片的尺寸修改，并转为tensor)
    image = transforms.Resize(target_size)(image)
    image = transforms.ToTensor()(image)

    # Convert to Torch.Tensor and normalize. mean与std   （RGB三通道）这里的参数和数据集中是对应的，训练过程中一致
    image = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(image)

    # Add batch_size axis.增加一个维度，用于按batch测试   本次这里一次测试一张
    image = image[None]
    if use_gpu:
        image = image.cuda()
    return Variable(image, volatile=True) #不需要求导

# 开启服务   这里的predict只是一个名字，可自定义
@app.route("/predict", methods=["POST"])
def predict():
    # Initialize the data dictionary that will be returned from the view.
    #做一个标志，刚开始无图像传入时为false，传入图像时为true
    data = {"success": False}

    # 如果收到请求
    if flask.request.method == 'POST':
        #判断是否为图像
        if flask.request.files.get("image"):
            # Read the image in PIL format
            # 将收到的图像进行读取
            image = flask.request.files["image"].read()
            image = Image.open(io.BytesIO(image)) #二进制数据

            # 利用上面的预处理函数将读入的图像进行预处理
            image = prepare_image(image, target_size=(64, 64))

            preds = F.softmax(model(image), dim=1)
            results = torch.topk(preds.cpu().data, k=3, dim=1)
            results = (results[0].cpu().numpy(), results[1].cpu().numpy())

            #将data字典增加一个key,value,其中value为list格式
            data['predictions'] = list()

            # Loop over the results and add them to the list of returned predictions
            for prob, label in zip(results[0][0], results[1][0]):
                #label_name = idx2label[str(label)]
                r = {"label": str(label), "probability": float(prob)}
                #将预测结果添加至data字典
                data['predictions'].append(r)

            # Indicate that the request was a success.
            data["success"] = True
    # 将最终结果以json格式文件传出
    return flask.jsonify(data)

"""
test_json = {
                "status_code": 200,
                "success": {
                            "message": "image uploaded",
                            "code": 200
                        },
                "video":{
                    "video_name":opt['source'].split('/')[-1],
                    "video_path":opt['source'],
                    "description":"1",
                    "length": str(hour)+','+str(minute)+','+str(round(second,4)),
                    "model_object_completed":model_flag
                    }
                    "status_txt": "OK"
                    }
                    response = requests.post(
                        'http://xxx.xxx.xxx.xxx:8090/api/ObjectToKafka/',,
                        data={'json': str(test_json)})
"""

if __name__ == '__main__':
    print("Loading PyTorch model and Flask starting server ...")
    print("Please wait until server has fully started")
    #先加载模型
    load_model()
    #再开启服务
    app.run(port='5012')

代码实现了一个基于 Flask 的图像分类预测服务，借助预训练的 PyTorch 模型对输入的图像进行分类预测。下面逐部分解释代码：

1. 导入模块

python 复制代码 
import io
import json
import flask
import torch
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
#from torchvision import transforms as T
from torchvision import transforms, models, datasets
from torch.autograd import Variable
  • io：用于处理二进制数据，这里用于读取图像的二进制流。
  • json：用于处理 JSON 数据，服务返回的预测结果会以 JSON 格式输出。
  • flask：用于创建 Web 服务。
  • torch：PyTorch 深度学习框架的核心模块。
  • torch.nn.functional：包含许多神经网络的函数，如 softmax 函数。
  • PIL.Image：用于处理图像，如读取、转换图像格式等。
  • torch.nn：用于构建神经网络模型。
  • torchvision.transforms：用于图像的预处理操作。
  • torchvision.models：提供预训练的模型，如 ResNet。
  • torch.autograd.Variable：用于封装张量，支持自动求导，不过在 PyTorch 0.4.0 之后，Tensor 已经集成了 Variable 的功能。

2. 初始化 Flask 应用和模型

python 复制代码 
# 初始化Flask app
app = flask.Flask(__name__)
model = None
use_gpu = False
  • app：Flask 应用实例。
  • model：全局变量，用于存储加载的模型。
  • use_gpu：布尔值，指示是否使用 GPU 进行推理。

3. 加载模型

python 复制代码 
def load_model():
    """Load the pre-trained model, you can use your model just as easily.
    """
    global model
    #这里我们直接加载官方工具包里提供的训练好的模型（代码会自动下载）括号内参数为是否下载模型对应的配置信息
    model = models.resnet18()
    num_ftrs = model.fc.in_features
    model.fc = nn.Sequential(nn.Linear(num_ftrs, 102))  # 类别数自己根据自己任务来

    #print(model)
    checkpoint = torch.load('best.pth')
    model.load_state_dict(checkpoint['state_dict'])
    #将模型指定为测试格式
    model.eval()
    #是否使用gpu
    if use_gpu:
        model.cuda()
  • models.resnet18()：加载预训练的 ResNet-18 模型。
  • model.fc：修改模型的全连接层，将输出类别数设置为 102。
  • torch.load('best.pth')：加载保存的模型参数。
  • model.load_state_dict(checkpoint['state_dict'])：将加载的参数应用到模型中。
  • model.eval()：将模型设置为评估模式，关闭一些在训练时使用的特殊层，如 Dropout。
  • model.cuda()：如果 use_gpu 为 True，将模型移动到 GPU 上。

4. 数据预处理

python 复制代码 
def prepare_image(image, target_size):
    """Do image preprocessing before prediction on any data.

    :param image:       original image
    :param target_size: target image size
    :return:
                        preprocessed image
    """
    #针对不同模型，image的格式不同，但需要统一至RGB格式
    if image.mode != 'RGB':
        image = image.convert("RGB")

    # Resize the input image and preprocess it.(按照所使用的模型将输入图片的尺寸修改，并转为tensor)
    image = transforms.Resize(target_size)(image)
    image = transforms.ToTensor()(image)

    # Convert to Torch.Tensor and normalize. mean与std   （RGB三通道）这里的参数和数据集中是对应的，训练过程中一致
    image = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(image)

    # Add batch_size axis.增加一个维度，用于按batch测试   本次这里一次测试一张
    image = image[None]
    if use_gpu:
        image = image.cuda()
    return Variable(image, volatile=True) #不需要求导
  • image.convert("RGB")：将图像转换为 RGB 格式。
  • transforms.Resize(target_size)：将图像调整为指定大小。
  • transforms.ToTensor()：将图像转换为 PyTorch 张量。
  • transforms.Normalize：对图像进行归一化处理。
  • image[None]：增加一个维度，模拟批量输入。
  • image.cuda()：如果 use_gpu 为 True，将图像移动到 GPU 上。
  • Variable(image, volatile=True)：封装张量，volatile=True 表示不需要计算梯度。

5. 定义预测接口

python 复制代码 
@app.route("/predict", methods=["POST"])
def predict():
    # Initialize the data dictionary that will be returned from the view.
    #做一个标志，刚开始无图像传入时为false，传入图像时为true
    data = {"success": False}

    # 如果收到请求
    if flask.request.method == 'POST':
        #判断是否为图像
        if flask.request.files.get("image"):
            # Read the image in PIL format
            # 将收到的图像进行读取
            image = flask.request.files["image"].read()
            image = Image.open(io.BytesIO(image)) #二进制数据

            # 利用上面的预处理函数将读入的图像进行预处理
            image = prepare_image(image, target_size=(64, 64))

            preds = F.softmax(model(image), dim=1)
            results = torch.topk(preds.cpu().data, k=3, dim=1)
            results = (results[0].cpu().numpy(), results[1].cpu().numpy())

            #将data字典增加一个key,value,其中value为list格式
            data['predictions'] = list()

            # Loop over the results and add them to the list of returned predictions
            for prob, label in zip(results[0][0], results[1][0]):
                #label_name = idx2label[str(label)]
                r = {"label": str(label), "probability": float(prob)}
                #将预测结果添加至data字典
                data['predictions'].append(r)

            # Indicate that the request was a success.
            data["success"] = True
    # 将最终结果以json格式文件传出
    return flask.jsonify(data)
  • @app.route("/predict", methods=["POST"])：定义一个 POST 请求的路由 /predict。
  • flask.request.files["image"]：获取上传的图像文件。
  • prepare_image：对图像进行预处理。
  • F.softmax(model(image), dim=1)：对模型的输出进行 softmax 处理，得到概率分布。
  • torch.topk(preds.cpu().data, k=3, dim=1)：获取概率最高的前 3 个类别及其概率。
  • flask.jsonify(data)：将预测结果以 JSON 格式返回。

6. 主程序

python 复制代码 
if __name__ == '__main__':
    print("Loading PyTorch model and Flask starting server ...")
    print("Please wait until server has fully started")
    #先加载模型
    load_model()
    #再开启服务
    app.run(port='5012')
  • load_model()：加载预训练模型。
  • app.run(port='5012')：启动 Flask 服务，监听端口 5012。

代码优化建议

  • 可以移除重复的 import torch 语句。
  • Variable 已经在 PyTorch 0.4.0 之后被弃用，可以直接使用 Tensor。
  • 可以添加异常处理，提高代码的健壮性。

2

python 复制代码 
import requests
import argparse

# url和端口携程自己的
flask_url = 'http://127.0.0.1:5012/predict'


def predict_result(image_path):
    #啥方法都行
    image = open(image_path, 'rb').read()
    payload = {'image': image}
    #request发给server.
    r = requests.post(flask_url, files=payload).json()

    # 成功的话在返回.
    if r['success']:
        # 输出结果.
        for (i, result) in enumerate(r['predictions']):
            print('{}. {}: {:.4f}'.format(i + 1, result['label'],
                                          result['probability']))
    # 失败了就打印.
    else:
        print('Request failed')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Classification demo')
    parser.add_argument('--file', default='./flower_data/train_filelist/image_06998.jpg', type=str, help='test image file')

    args = parser.parse_args()
    predict_result(args.file)

代码的主要功能是向一个基于 Flask 搭建的图像分类预测服务发送请求，并处理返回的预测结果。下面逐部分解释代码：

1. 导入模块

python 复制代码 
import requests
import argparse
  • requests：这是一个常用的 Python HTTP 库，用于发送 HTTP 请求，在代码里主要用于向 Flask 服务发送 POST 请求。
  • argparse：Python 标准库中的命令行参数解析模块，可让程序在运行时从命令行接收参数。

2. 定义 Flask 服务的 URL

python 复制代码 
# url和端口写成自己的
flask_url = 'http://127.0.0.1:5012/predict'

这里定义了要请求的 Flask 服务的 URL，http://127.0.0.1:5012 表示本地主机的 5012 端口，/predict 是服务中处理预测请求的路由。

3. 定义预测函数

python 复制代码 
def predict_result(image_path):
    # 啥方法都行
    image = open(image_path, 'rb').read()
    payload = {'image': image}
    # request发给server.
    r = requests.post(flask_url, files=payload).json()

    # 成功的话就返回.
    if r['success']:
        # 输出结果.
        for (i, result) in enumerate(r['predictions']):
            print('{}. {}: {:.4f}'.format(i + 1, result['label'],
                                          result['probability']))
    # 失败了就打印.
    else:
        print('Request failed')
  • image = open(image_path, 'rb').read()：以二进制只读模式打开指定路径的图像文件，并读取其内容。
  • payload = {'image': image}：创建一个字典 payload，将读取的图像数据作为值，键为 'image'。
  • r = requests.post(flask_url, files=payload).json()：使用 requests.post 方法向指定的 Flask 服务 URL 发送 POST 请求，将图像数据作为文件上传。json() 方法将服务返回的 JSON 数据解析为 Python 字典。
  • if r['success']：检查服务返回的结果中 success 字段是否为 True，如果是，则遍历 predictions 列表并打印每个预测结果的排名、类别标签和概率；否则，打印请求失败的信息。

4. 主程序部分

python 复制代码 
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Classification demo')
    parser.add_argument('--file', default='./flower_data/train_filelist/image_06998.jpg', type=str, help='test image file')

    args = parser.parse_args()
    predict_result(args.file)
  • argparse.ArgumentParser：创建一个命令行参数解析器对象，设置描述信息为 'Classification demo'。
  • parser.add_argument：添加一个名为 --file 的命令行参数，默认值为 './flower_data/train_filelist/image_06998.jpg'，类型为字符串，用于指定要测试的图像文件路径。
  • args = parser.parse_args()：解析命令行参数，并将结果存储在 args 对象中。
  • predict_result(args.file)：调用 predict_result 函数，传入命令行指定的图像文件路径进行预测。

总结

代码实现了一个简单的客户端程序，用于向 Flask 服务发送图像分类预测请求，并处理和输出预测结果。用户可以通过命令行指定要测试的图像文件路径。

3 模型预测结果

98号的概率比较高，能够实现一定的预测效果

相关推荐
Hali_Botebie35 分钟前
【蒸馏用损失】NCEloss介绍，大规模分类任务的损失函数
人工智能·分类·数据挖掘
阿里云云原生35 分钟前
RAG 调优指南：Spring AI Alibaba 模块化 RAG 原理与使用
java·人工智能·spring
niuniu_66636 分钟前
selenium应用测试场景
python·selenium·测试工具·单元测试·测试
蚍蜉撼树谈何易38 分钟前
机器学习的定义及分类
人工智能·机器学习·分类
城电科技2 小时前
城电科技｜零碳美丽示范村建设方案 能源+景观+教育
人工智能·科技·生活·能源
满怀10152 小时前
Python扩展知识详解：lambda函数
开发语言·python
佚名涙3 小时前
go中锁的入门到进阶使用
开发语言·后端·golang
蓝博AI4 小时前
基于卷积神经网络的眼疾识别系统，resnet50，efficentnet（pytorch框架，python代码）
pytorch·python·cnn
牧歌悠悠5 小时前
【Python 算法】动态规划
python·算法·动态规划
opentrending6 小时前
Github 热点项目 awesome-mcp-servers MCP 服务器合集，3分钟实现AI模型自由操控万物！
服务器·人工智能·github
热门推荐
01我决定放弃搞 Java 了02DeepSeek各版本说明与优缺点分析03如何在WPS和Word/Excel中直接使用DeepSeek功能04从零安装 LLaMA-Factory 微调 Qwen 大模型成功及所有的坑05苍穹外卖面试总结06DeepSeek R1本地化部署 Ollama + Chatbox 打造最强 AI 工具07论文笔记：交替单模态适应的多模态表征学习08RAG 实践- Ollama+RagFlow 部署本地知识库09如何本地部署AI智能体平台，带你手搓一个AI Agent10三次握手内部实现原理