深度学习模型部署篇——从0部署深度学习分类模型（二）-灵析社区

使用模型部署-->预测摄像头实时画面

其实不管是预测一张图像还是实时对视频进行预测，原理都是差不多的，我们一起来看一下。🧃🧃🧃

首先，导入一些必要的工具包：

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.pth的权重文件，这个是我们部署模型的关键。

有了AlexNet.pth权重文件，我们要对其进行加载，对模型保存和加载不清楚的请阅读下方博客：