图像分类识别——（4）分类识别视频文件

需要对视频文件逐帧的分割后进行图像分类，然后再把分割后的图片重新串成一个新视频。

 

导入工具包

# 导入工具包，操作系统相关的包
import os
import time
import shutil
import tempfile
from tqdm import tqdm

# 导入图像处理相关的包
import cv2
from PIL import Image

# 导入与可视化相关的包
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline
plt.rcParams['axes.unicode_minus']=False # 用来正常显示负号
plt.rcParams['font.sans-serif']=['SimHei'] # 用来正常显示中文标签
import gc

# 导入pytorch相关的包
import torch
import torch.nn.functional as F
from torchvision import models

# 运行前需保证已经按照mmcv库，如没有安装，看之前环境配置的文件

import mmcv

# 获取设备
# 有 GPU 就用 GPU，没有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device:', device)

 

#这次要生成的是视频文件，所以不需要显示每一帧的绘图
# 后端绘图，不显示，只保存
import matplotlib
matplotlib.use('Agg')

 

载入预训练图像分类模型

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['wordnet'], row['class']]

idx_to_labels

 

图像预处理

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 pred_single_frame(img, n=5):
'''
输入摄像头画面bgr-array，输出前n个图像分类预测结果的图像bgr-array
'''
img_bgr = img
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, n) # 取置信度最大的 n 个结果
pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析出类别
confs = top_n[0].cpu().detach().numpy().squeeze() # 解析出置信度

# 在图像上写字
for i in range(n):
class_name = idx_to_labels[pred_ids[i]][1] # 获取类别名称
confidence = confs[i] * 100 # 获取置信度
text = '{:<15} {:>.4f}'.format(class_name, confidence)

# !图片，添加的文字，左上角坐标，字体，字号，bgr颜色，线宽
img_bgr = cv2.putText(img_bgr, text, (25, 50 + 40 * i), cv2.FONT_HERSHEY_SIMPLEX, 1.25, (0, 0, 255), 3)

return img_bgr, pred_softmax

 

视频处理

# 输入输出视频路径
input_video = 'test_img/video_3.mp4'

# 创建临时文件夹，存放每帧结果,文件名是现在的时间。用来存储每一帧的预测结果
temp_out_dir = time.strftime('%Y%m%d%H%M%S')
os.mkdir(temp_out_dir)
print('创建文件夹 {} 用于存放每帧预测结果'.format(temp_out_dir))

 

# 读入待预测视频
imgs = mmcv.VideoReader(input_video)

prog_bar = mmcv.ProgressBar(len(imgs))

# 对视频逐帧处理
for frame_id, img in enumerate(imgs):

## 处理单帧画面
img, pred_softmax = pred_single_frame(img, n=5)

# 将处理后的该帧画面图像文件，保存至 /tmp 目录下
cv2.imwrite(f'{temp_out_dir}/{frame_id:06d}.jpg', img)

prog_bar.update() # 更新进度条

# 把每一帧串成视频文件
mmcv.frames2video(temp_out_dir, 'output/output_pred.mp4', fps=imgs.fps, fourcc='mp4v')

shutil.rmtree(temp_out_dir) # 删除存放每帧画面的临时文件夹
print('删除临时文件夹', temp_out_dir)