本文为稀土掘金技术社区首发签约文章，30天内禁止转载，30天后未获授权禁止转载，侵权必究！
🍊作者简介：秃头小苏，致力于用最通俗的语言描述问题
🍊专栏推荐：深度学习网络原理与实战
🍊近期目标：写好专栏的每一篇文章
🍊支持小苏：点赞👍🏼、收藏⭐、留言📩
深度学习模型部署篇——从0部署深度学习分类模型（二）
写在前面
Hello，大家好，我是小苏👦🏽👦🏽👦🏽
在上一节，我已经为大家介绍了什么是模型部署，为什么要进行模型部署已经模型部署的流程，大家消化了多少呢。🥗🥗🥗如果对模型部署的流程还不熟悉的话建议去看一下上期博客。
上节末尾也和大家提到，虽然我们说了模型部署可以加速我们的推理速度，让模型能够应用于生产环境当中。但是上节我们只对单张图像进行了推理，并不能很直观的感受推理速度的变化。那么这节我将调用PC端的摄像头来让大家体验一下🍉🍉🍉
准备好了的话，我们就起飞咯🚀🚀🚀
使用模型部署–>预测摄像头实时画面
其实不管是预测一张图像还是实时对视频进行预测，原理都是差不多的，我们一起来看一下。🧃🧃🧃
首先，导入一些必要的工具包：
import os
import cv2
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from torchvision import transforms
import onnxruntime
from PIL import Image, ImageFont, ImageDraw
import matplotlib.pyplot as plt
因为我们想要在图像上展示中文，需要下载中文字体并加载：
# 下载中文字体文件
wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/dataset/SimHei.ttf
# 导入中文字体，指定字体大小
font = ImageFont.truetype('SimHei.ttf', 32)
接着我们载入ONNX和ImageNet1000的分类标签，这些和上一节都是一致的：
# 载入 onnx 模型
ort_session = onnxruntime.InferenceSession('resnet18_imagenet.onnx')
# 载入ImageNet 1000图像分类标签
df = pd.read_csv('imagenet_class_index.csv')
idx_to_labels = {}
for idx, row in df.iterrows():
idx_to_labels[row['ID']] = row['Chinese']
同时，也定义一个预处理的函数，和第一节是一致的：
# 测试集图像预处理-RCTN：缩放裁剪、转 Tensor、归一化
test_transform = transforms.Compose([transforms.Resize(256),
transforms.CenterCrop(256),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
这些算是准备工作，然后我们来思考一下如何来做摄像头的实时预测。其实很容易啦，我们会一帧一帧的处理视频中的图像，而一帧图像的处理过程是不是就和我们上一节讲的一致呢，为了方便调用，我们把处理一帧图像的方法封装起来，如下：
# 处理一帧图像的函数
def process_frame(img_bgr):
# 记录该帧开始处理的时间
start_time = time.time()
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) # BGR转RGB
img_pil = Image.fromarray(img_rgb) # array 转 PIL
## 预处理
input_img = test_transform(img_pil) # 预处理
input_tensor = input_img.unsqueeze(0).numpy()
## onnx runtime 预测
ort_inputs = {'input': input_tensor} # onnx runtime 输入
pred_logits = ort_session.run(['output'], ort_inputs)[0] # onnx runtime 输出
pred_logits = torch.tensor(pred_logits)
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算
## 解析图像分类预测结果
n = 5
top_n = torch.topk(pred_softmax, n) # 取置信度最大的 n 个结果
pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析出类别
confs = top_n[0].cpu().detach().numpy().squeeze() # 解析出置信度
## 在图像上写中文
draw = ImageDraw.Draw(img_pil)
for i in range(len(confs)):
pred_class = idx_to_labels[pred_ids[i]]
# 写中文：文字坐标，中文字符串，字体，rgba颜色
text = '{:<15} {:>.3f}'.format(pred_class, confs[i]) # 中文字符串
draw.text((50, 100 + 50 * i), text, font=font, fill=(255, 0, 0, 1))
img_rgb = np.array(img_pil) # PIL 转 array
img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) # RGB转BGR
# 记录该帧处理完毕的时间
end_time = time.time()
# 计算每秒处理图像帧数FPS
FPS = 1/(end_time - start_time)
# 图片，添加的文字，左上角坐标，字体，字体大小，颜色，线宽，线型
img_bgr = cv2.putText(img_bgr, 'FPS '+str(int(FPS)), (50, 80), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 255), 4, cv2.LINE_AA)
return img_bgr
这个函数和上节的过程基本是一致的，不过这里我们添加了计算FPS的代码，并将其显示在图像上：
start_time = time.time() # 图像处理开始
"""
---图像处理过程中
---图像处理过程中
---图像处理过程中
"""
end_time = time.time() # 图像处理结束
# 计算FPS
FPS = 1/(end_time - start_time)
FPS表示一秒钟处理多少帧的图像，比如(end_time - start_time)=0.5 ，则FPS=1/0.5=2，表示每秒钟能处理两帧图像，显然FPS值越大，表示处理速度越快。🍄🍄🍄
定义好处理一帧图像的函数，我们就可以调用摄像头进行实时预测了：
import cv2
import time
# 获取摄像头，传入0表示获取系统默认摄像头
cap = cv2.VideoCapture(0)
# 打开cap
cap.open(0)
# 无限循环，直到break被触发
while cap.isOpened():
# 获取画面
success, frame = cap.read()
if not success: # 如果获取画面不成功，则退出
print('获取画面不成功，退出')
break
## 逐帧处理
frame = process_frame(frame)
# 展示处理后的三通道图像
cv2.imshow('my_window',frame)
key_pressed = cv2.waitKey(60) # 每隔多少毫秒毫秒，获取键盘哪个键被按下
# print('键盘上被按下的键：', key_pressed)
if key_pressed in [ord('q'),27]: # 按键盘上的q或esc退出（在英文输入法下）
break
# 关闭摄像头
cap.release()
# 关闭图像窗口
cv2.destroyAllWindows()
我们可以来看一下我们处理的结果，如下图所示：
可以看到，模型可以检测到我们的物体。当然了，我的背景比较杂，所以识别率并没有很高。但是本节我也不关注识别率，让我们来看看FPS是多少，可以看到，大概是18-20左右。🍦🍦🍦
不使用模型部署–>预测摄像头实时画面
上节展示了使用模型部署来预测摄像头实时画面的案例，FPS大概在18左右，这次我们不使用模型部署来看看FPS大概是多少。
可以看到不使用模型部署时，FPS只在9左右，前前后后相差了一倍之多，所以说在工业实践中模型部署还是非常有必要的。
我也贴出这部分的代码叭，如下：
import numpy as np
import pandas as pd
from PIL import Image, ImageFont, ImageDraw
import cv2
import time
import torch
import torch.nn.functional as F
from torchvision import models
# 有 GPU 就用 GPU，没有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device:', device)
# 导入中文字体，指定字号
font = ImageFont.truetype('SimHei.ttf', 32)
model = models.resnet18(pretrained=True)
model = model.eval()
model = model.to(device)
# 载入ImageNet 1000图像分类标签
df = pd.read_csv('imagenet_class_index.csv')
idx_to_labels = {}
for idx, row in df.iterrows():
idx_to_labels[row['ID']] = row['Chinese']
from torchvision import transforms
# 测试集图像预处理-RCTN：缩放裁剪、转 Tensor、归一化
test_transform = 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 process_frame(img):
# 记录该帧开始处理的时间
start_time = time.time()
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # BGR转RGB
img_pil = Image.fromarray(img_rgb) # array 转 PIL
input_img = test_transform(img_pil).unsqueeze(0).to(device) # 预处理
pred_logits = model(input_img) # 执行前向预测，得到所有类别的 logit 预测分数
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算
top_n = torch.topk(pred_softmax, 5) # 取置信度最大的 n 个结果
pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析预测类别
confs = top_n[0].cpu().detach().numpy().squeeze() # 解析置信度
# 使用PIL绘制中文
draw = ImageDraw.Draw(img_pil)
# 在图像上写字
for i in range(len(confs)):
pred_class = idx_to_labels[pred_ids[i]]
text = '{:<15} {:>.3f}'.format(pred_class, confs[i])
# 文字坐标，中文字符串，字体，bgra颜色
draw.text((50, 100 + 50 * i), text, font=font, fill=(255, 0, 0, 1))
img = np.array(img_pil) # PIL 转 array
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) # RGB转BGR
# 记录该帧处理完毕的时间
end_time = time.time()
# 计算每秒处理图像帧数FPS
FPS = 1 / (end_time - start_time)
# 图片，添加的文字，左上角坐标，字体，字体大小，颜色，线宽，线型
img = cv2.putText(img, 'FPS ' + str(int(FPS)), (50, 80), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 255), 4,
cv2.LINE_AA)
return img
#获取摄像头，传入0表示获取系统默认摄像头
cap = cv2.VideoCapture(1)
# 打开cap
cap.open(0)
# 无限循环，直到break被触发
while cap.isOpened():
# 获取画面
success, frame = cap.read()
if not success:
print('Error')
break
## !!!处理帧函数
frame = process_frame(frame)
# 展示处理后的三通道图像
cv2.imshow('my_window', frame)
if cv2.waitKey(1) in [ord('q'), 27]: # 按键盘上的q或esc退出（在英文输入法下）
break
# 关闭摄像头
cap.release()
# 关闭图像窗口
cv2.destroyAllWindows()
使用模型部署–>预测视频文件
前面两个小节是调用PC的摄像头进行预测，但是有时候我们会对视频进行预测，这该怎么做呢，我们一起来看看叭。🌱🌱🌱
这里我只展示处理视频的函数，一些包和前面的准备工作就不在赘述了：
-
定义
generate_video函数，并设置input_path参数，默认为 “videos.mp4″。该函数用于处理视频。def generate_video(input_path='videos.mp4'):定义一个函数用于处理视频，后面直接调用即可，调用代码为：
generate_video(input_path='video_4.mp4') -
解析输入视频的文件名，并生成输出视频的路径。打印提示信息。
filehead = input_path.split('/')[-1] output_path = "out-" + filehead print('视频开始处理',input_path) -
使用
cv2.VideoCapture打开视频文件，并计算视频的总帧数。这里使用一个循环遍历获取每一帧，计算帧数，并在循环结束后释放资源。cap = cv2.VideoCapture(input_path) frame_count = 0 while(cap.isOpened()): success, frame = cap.read() frame_count += 1 if not success: break cap.release() print('视频总帧数为',frame_count) -
再次使用
cv2.VideoCapture打开视频文件，并获取视频的帧大小。cap = cv2.VideoCapture(input_path) frame_size = (cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) -
创建一个
cv2.VideoWriter对象out，用于将处理后的帧写入输出视频文件。设置视频编码器为 “mp4v”，帧率为输入视频的帧率。此处还可以根据需要设置输出视频的参数。fourcc = cv2.VideoWriter_fourcc(*'mp4v') fps = cap.get(cv2.CAP_PROP_FPS) out = cv2.VideoWriter(output_path, fourcc, fps, (int(frame_size[0]), int(frame_size[1]))) -
使用
tqdm库创建一个进度条对象pbar，并将总帧数设置为进度条的总量。with tqdm(total=frame_count) as pbar: # 进入主循环，遍历每一帧图像。在每次迭代中，使用 cap.read() 读取视频的下一帧图像，并检查是否成功读取。 while(cap.isOpened()): success, frame = cap.read() if not success: break -
调用
process_frame函数对每一帧图像进行处理。如果处理过程中出现错误，会打印错误信息。try: frame = process_frame(frame) except Exception as error: print('报错！', error) pass -
如果读取和处理图像成功，将处理后的帧图像写入输出视频文件。
if success == True: out.write(frame) -
更新进度条的进度。
pbar.update(1) -
循环结束后，释放所有使用的资源，关闭窗口，并保存输出视频。
cv2.destroyAllWindows() out.release() cap.release() print('视频已保存', output_path)
让我们来看看结果叭。🍮🍮🍮
使用模型部署–>部署自定义的花5分类模型
前文讲述的都是我们用官方训练好的模型进行部署，那么针对我们自己的数据集，我们训练了一个分类模型，我们应该如何进行部署呢？让我们一起来看看叭~~~🍄🍄🍄
在之前的博客中，我们一直用花的五分类来做物体分类实验，那么这里我们同样来实现花的五分类模型部署。因为是教学，我们直接使用最简单的AlexNet模型实现花的五分类，对此不熟悉的请先阅读下方博客：
看完上述博客，我们知道我们训练得到了名为AlexNet.pth的权重文件，这个是我们部署模型的关键。
有了AlexNet.pth权重文件，我们要对其进行加载，对模型保存和加载不清楚的请阅读下方博客：
我们使用的保存模型的方式是通过官方推荐的方式一进行保存的，因此我们要通过方式一来加载模型。
-
实列化AlexNet模型
# 搭建网络模型 class AlexNet(nn.Module): def __init__(self, num_classes=5): super(AlexNet, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 96, 11, 4, padding=0), nn.ReLU(inplace=True), nn.MaxPool2d(3, 2, padding=0), nn.Conv2d(96, 256, 5, padding=2), nn.ReLU(inplace=True), nn.MaxPool2d(3, 2, padding=0), nn.Conv2d(256, 384, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(384, 384, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(384, 256, 3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(3, 2, padding=0), ) self.classifier = nn.Sequential( nn.Flatten(), nn.Dropout(), nn.Linear(256 * 6 * 6, 4096), nn.ReLU(inplace=True), nn.Dropout(), nn.Linear(4096, 4096), nn.ReLU(inplace=True), nn.Linear(4096, num_classes) ) def forward(self, x): x = self.features(x) x = self.classifier(x) return x model = AlexNet(num_classes=5).to(device) -
加载模型并设置为推理模式
# 加载模型 model.load_state_dict(torch.load('checkpoint/AlexNet.pth', map_location='cpu')) # 设置为推理模式 model = model.eval().to(device)
模型加载完成后，我们将模型转成ONNX中间格式：
x = torch.randn(1, 3, 255, 227).to(device)
output = model(x)
with torch.no_grad():
torch.onnx.export(
model, # 要转换的模型
x, # 模型的任意一组输入
'Alex_flower5.onnx', # 导出的 ONNX 文件名
opset_version=11, # ONNX 算子集版本
input_names=['input'], # 输入 Tensor 的名称（自己起名字）
output_names=['output'] # 输出 Tensor 的名称（自己起名字）
)
本地产生ONNX格式文件：
双击查看一下：
到这里我们已经得到了花的五分类的ONNX中间格式，后面的操作就和前一节一样了，即加载ONNX模型并通过ONNX Runtime推理引擎来进行模型推理，这里我就不过多叙述了。我从网上下载的推理图像如下【这个是郁金香花🌷🌷🌷】：
推理完成后我们同样可以得到预测类别pred_ids和预测置信度confs，她们的值如下：
我们可以载入类别和对应ID展示一下：
# 载入类别和对应 ID
idx_to_labels = np.load('idx_to_labels1.npy', allow_pickle=True).item()
for i in range(n):
class_name = idx_to_labels[pred_ids[i]] # 获取类别名称
confidence = confs[i] * 100 # 获取置信度
text = '{:<6} {:>.3f}'.format(class_name, confidence)
print(text)
其中，idx_to_labels1.npy的文件内容如下：
最终输出的结果如下，可以看出，正确的预测出来输入图片为郁金香。
小结
好啦，这节就为大家介绍到这里咯。🌷🌷🌷大家快去试试叭，看看自己模型的推理速度有没有嗖嗖的增加喔。🥂🥂🥂
参考文献
咳咳，重点来了。今天可是七夕哎，希望大家都能够拥有属于自己的幸福喔🍄🍄🍄
如若文章对你有所帮助，那就🛴🛴🛴
