文章目录PyTorch 图像分类预测脚本完整代码一览导入依赖库与全局变量模型加载函数 load_model图像预处理函数 prepare_image预测函数 predict主程序入口Flask Web 服务封装完整代码一览新增导入模块创建 Flask 应用实例复用 load_model 和 prepare_image定义路由 /predict启动 Web 服务测试你的 Web 服务PyTorch 图像分类预测脚本完整代码一览import torch import torch.nn.functional as F from PIL import Image from torch import nn from torchvision import transforms,models modelNone use_gpuFalse defload_model():加载预训练模型global model modelmodels.resnet18()num_ftrsmodel.fc.in_features model.fcnn.Sequential(nn.Linear(num_ftrs,out_features102))checkpointtorch.load(best.pth)model.load_state_dict(checkpoint[state_dict])model.eval()ifuse_gpu:model.cuda()defprepare_image(image,target_size):ifimage.mode!RGB:imageimage.convert(RGB)imagetransforms.Resize(target_size)(image)imagetransforms.ToTensor()(image)imagetransforms.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225])(image)imageimage[None]# 增加 batch 维度ifuse_gpu:imageimage.cuda()returnimage defpredict(image):data{success:False}imageprepare_image(image,target_size(224,224))predsF.softmax(model(image),dim1)resultstorch.topk(preds.cpu(),k3,dim1)results(results[0].detach().cpu().numpy(),results[1].detach().cpu().numpy())data[predictions]list()forprob,label inzip(results[0][0],results[1][0]):r{label:str(label),probability:float(prob)}data[predictions].append(r)data[success]Truereturndataif__name____main__:load_model()img_pathr./flower_data/train_filelist/image_00014.jpgtest_imgImage.open(img_path)respredict(test_img)print(res)导入依赖库与全局变量import torch import torch.nn.functional as F from PIL import Image from torch import nn from torchvision import transforms,models modelNone use_gpuFalsemodel None全局变量用于存放加载后的模型。use_gpu False是否启用 GPU 加速这里设为 False如需启用需自行安装 CUDA 环境并设为 True。模型加载函数 load_modeldefload_model():global model modelmodels.resnet18()# 创建一个 ResNet18 模型实例使用 ImageNet 预训练权重 num_ftrsmodel.fc.in_features # ResNet18 的全连接层fc的输入特征数默认为512 model.fcnn.Sequential(nn.Linear(num_ftrs,out_features102))# 将原来的全连接层替换为一个新的线性层输出为102个类别 checkpointtorch.load(best.pth)# 加载训练好的模型权重文件 model.load_state_dict(checkpoint[state_dict])# 将权重加载到模型中 model.eval()# 将模型设置为评估模式ifuse_gpu:model.cuda()图像预处理函数 prepare_imagedefprepare_image(image,target_size):ifimage.mode!RGB:# 如果图像模式不是 RGB比如 RGBA 或 L先转换为 RGB确保三通道输入 imageimage.convert(RGB)imagetransforms.Resize(target_size)(image)# 将图像缩放到(224,224)这是 ResNet 的标准输入尺寸 imagetransforms.ToTensor()(image)# 将 PIL 图像或 NumPy 数组转换为 PyTorch 张量 imagetransforms.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225])(image)# 按均值和标准差进行标准化使数据分布接近标准正态分布 imageimage[None]# 增加 batch 维度ifuse_gpu:imageimage.cuda()returnimage # 返回预处理后的张量预测函数 predictdefpredict(image):data{success:False}#创建一个字典 data包含success键默认为 False用于表示是否成功完成预测 imageprepare_image(image,target_size(224,224))#调用 prepare_image 预处理图片 predsF.softmax(model(image),dim1)#在类别维度dim1上计算 softmax将 logits 转换为概率 resultstorch.topk(preds.cpu(),k3,dim1)#返回概率最大的前3个类别及其索引 results(results[0].detach().cpu().numpy(),results[1].detach().cpu().numpy())data[predictions]list()forprob,label inzip(results[0][0],results[1][0]):r{label:str(label),probability:float(prob)}data[predictions].append(r)data[success]Truereturndata返回结果是一个元组 (values, indices)分别对应概率值和类别索引。我们将其转为 NumPy 数组方便遍历。遍历前 3 个结果构建一个列表每个元素是一个字典包含 “label”类别索引和 “probability”概率值。将列表赋值给 data[‘predictions’]并将 “success” 设为 True。返回 data 字典便于后续转为 JSON 格式进行网络传输。主程序入口if__name____main__:#保证只有在直接运行此脚本时才执行以下代码被导入时不执行load_model()img_pathr./flower_data/train_filelist/image_00014.jpgtest_imgImage.open(img_path)respredict(test_img)print(res)调用 load_model() 加载模型权重。指定一张测试图片的路径用 Image.open 读取图片PIL 格式然后调用 predict 函数得到预测结果。打印结果字典。运行结果{success:True,predictions:[{label:22,probability:0.8882589936256409},{label:91,probability:0.07465873658657074},{label:13,probability:0.016419561579823494}]}Flask Web 服务封装完整代码一览import io import flask import torch import torch.nn.functional as F from PIL import Image from torch import nn from torchvision import transforms,models appflask.Flask(__name__)modelNone use_gpuFalse defload_model():加载预训练模型global model modelmodels.resnet18()num_ftrsmodel.fc.in_features model.fcnn.Sequential(nn.Linear(num_ftrs,out_features102))checkpointtorch.load(best.pth)model.load_state_dict(checkpoint[state_dict])model.eval()ifuse_gpu:model.cuda()defprepare_image(image,target_size):ifimage.mode!RGB:imageimage.convert(RGB)imagetransforms.Resize(target_size)(image)imagetransforms.ToTensor()(image)imagetransforms.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225])(image)imageimage[None]ifuse_gpu:imageimage.cuda()returnimage app.route(/predict,methods[POST])defpredict():data{success:False}ifflask.request.methodPOST:ifflask.request.files.get(image):image_bytesflask.request.files[image].read()imageImage.open(io.BytesIO(image_bytes))imageprepare_image(image,target_size(224,224))predsF.softmax(model(image),dim1)resultstorch.topk(preds.cpu().data,k3,dim1)results(results[0].detach().cpu().numpy(),results[1].detach().cpu().numpy())data[predictions]list()forprob,label inzip(results[0][0],results[1][0]):r{label:str(label),probability:float(prob)}data[predictions].append(r)data[success]True response{message:success,code:200,data:data}returnflask.jsonify(response)returnflask.jsonify(data)if__name____main__:print(Loading PyTorch model and Flask starting server ...)print(Please wait until server has fully started)load_model()app.run(host0.0.0.0,port5012,debugTrue)新增导入模块import io import flaskioPython 标准库提供 BytesIO用于在内存中读写二进制数据。flaskWeb 框架需要提前安装pip install flask。创建 Flask 应用实例appflask.Flask(__name__)初始化一个 Flask 应用对象name告诉 Flask 当前模块的名称用于定位资源文件。复用 load_model 和 prepare_image这两个函数与第一份代码完全一样我们直接复制过来。这体现了代码复用思想核心逻辑不变只在外围增加 Web 交互层。定义路由 /predictapp.route(/predict,methods[POST])defpredict():data{success:False}ifflask.request.methodPOST:#首先检查请求方法是否为 POSTifflask.request.files.get(image):#然后检查 flask.request.files 中是否包含名为image的文件 image_bytesflask.request.files[image].read()#读取文件的二进制内容 imageImage.open(io.BytesIO(image_bytes))#将二进制字节包装成类文件对象供 PIL.Image.open 读取 imageprepare_image(image,target_size(224,224))#调用 prepare_image 进行预处理然后执行前向传播 #...预测逻辑与第一份相同...data[success]Truereturnflask.jsonify({message:success,code:200,data:data})#将结果存入 data 字典returnflask.jsonify(data)#将 Python 字典转为 JSON 格式的 HTTP 响应并设置正确的 Content-Typeapp.route(“/predict”, methods[“POST”]) 装饰器将下面定义的 predict 函数绑定到 URL /predict并限制只接受 POST 请求。启动 Web 服务if__name____main__:print(Loading PyTorch model and Flask starting server ...)load_model()app.run(host0.0.0.0,port5012,debugTrue)先调用 load_model() 加载模型服务启动时只加载一次所有请求共享。app.run() 启动 Flask 内置服务器host‘0.0.0.0’监听所有网络接口允许其他设备访问。port5012使用端口 5012。debugTrue开启调试模式代码修改后自动重启方便开发。运行结果*Serving Flask appFlaskweb*Debug mode:on WARNING:This is a development server.Do not use it in a production deployment.Use a production WSGI server instead.*Running on alladdresses(0.0.0.0)*Running on http://127.0.0.1:5012*Running on http://192.168.31.33:5012Press CTRLC to quit*Restarting with stat Loading PyTorch model and Flask starting server...Please wait until server has fully started D:\lanzhi_study\深度学习\flower_data\Flaskweb.py:25:FutureWarning:You are using torch.load with weights_onlyFalse(the currentdefaultvalue),which uses thedefaultpickle module implicitly.It is possible to construct malicious pickle data which will execute arbitrary code duringunpickling(See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for weights_only will be flipped to True. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via torch.serialization.add_safe_globals. We recommend you start setting weights_onlyTrue for any use case where you dont have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.checkpointtorch.load(best.pth)*Debugger is active!*Debugger PIN:949-937-993测试你的 Web 服务服务启动后可以用两种方式测试方式一使用 curl 命令命令行在 命令提示符cmd 或终端中执行curl-X POST-FimageD:\lanzhi_study\深度学习\flower_data\flower_data\train_filelist\image_00001.jpghttp://127.0.0.1:5012/predict-X POST指定请求方法为 POST。-F “image图片路径”以 form-data 方式上传文件 后面跟本地图片的绝对路径。http://127.0.0.1:5012/predict接口地址。你会收到 JSON 格式的预测结果。运行结果方式二使用 Postman图形化工具请求方式改为 POST。URLhttp://127.0.0.1:5012/predict。Body → form-dataKey输入 imageKey 右侧的下拉框选择 File不是 Text。Value点击 Select Files选择一张本地图片。点击 Send 按钮下方会返回预测结果。
》)
