1
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_featuresmodel.fc = nn.Sequential(nn.Linear(num_ftrs, 102)) # 类别数自己根据自己任务来#print(model)checkpoint = torch.load('best.pth')model.load_state_dict(checkpoint['state_dict'])#将模型指定为测试格式model.eval()#是否使用gpuif 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，传入图像时为truedata = {"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 predictionsfor 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. 导入模块
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 应用和模型
# 初始化Flask app
app = flask.Flask(__name__)
model = None
use_gpu = False
app：Flask 应用实例。model：全局变量，用于存储加载的模型。use_gpu：布尔值，指示是否使用 GPU 进行推理。
3. 加载模型
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_featuresmodel.fc = nn.Sequential(nn.Linear(num_ftrs, 102)) # 类别数自己根据自己任务来#print(model)checkpoint = torch.load('best.pth')model.load_state_dict(checkpoint['state_dict'])#将模型指定为测试格式model.eval()#是否使用gpuif 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. 数据预处理
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. 定义预测接口
@app.route("/predict", methods=["POST"])
def predict():# Initialize the data dictionary that will be returned from the view.#做一个标志，刚开始无图像传入时为false，传入图像时为truedata = {"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 predictionsfor 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. 主程序
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
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. 导入模块
import requests
import argparse
requests：这是一个常用的 Python HTTP 库，用于发送 HTTP 请求，在代码里主要用于向 Flask 服务发送 POST 请求。argparse：Python 标准库中的命令行参数解析模块，可让程序在运行时从命令行接收参数。
2. 定义 Flask 服务的 URL
# url和端口写成自己的
flask_url = 'http://127.0.0.1:5012/predict'
这里定义了要请求的 Flask 服务的 URL，http://127.0.0.1:5012 表示本地主机的 5012 端口，/predict 是服务中处理预测请求的路由。
3. 定义预测函数
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. 主程序部分
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号的概率比较高，能够实现一定的预测效果