RT-DETR 目标检测视频推理
先前介绍了RT-DETR推理单张图像的案例,今天则介绍以下如何利用RT-DETR来进行视频推理。
事实上,进行视频推理的过程与单张图片的过程及其类似,就是将原本的视频切分为多帧图像后再进行推理即可。这里面涉及到Image等相关操作,今天便借此机会梳理一遍。
我们的实现思路很简单,其步骤如下:
- 将视频拆分为多帧
- 将多帧图像依次输入RT-DETR模型中进行检测
- 将多帧图像的检测结果合并为视频
利用cv2生成视频读取器,读取视频
classes = ['car','truck',"bus"] videoname = '1.mp4' capture = cv2.VideoCapture(videoname) images = []
通过循环将视频拆分为多帧图像
if capture.isOpened(): while True: ret,img=capture.read() # img 就是一帧图片,此时的img是ndarray格式,即numpy的数组类型 if not ret:break # 当获取完最后一帧就结束 img=cv2.resize(img,(640,640))#由于模型输入的图像尺寸固定为640*640,因此需要转换 img = Image.fromarray(img)#将数组类型转换为Image类型 img = img.convert('RGB')#通道位置置换 im_data = ToTensor()(img)[None]#模型输入类型为tensor,故将其转换为tensor类型 size = torch.tensor([[640, 640]])#设定输入图像的尺寸
加载onnx模型文件,并判断GPU是否可用,同时进行前向推理,并计算FPS
if torch.cuda.is_available(): print("GPU") sess = ort.InferenceSession("model.onnx", None, providers=["CUDAExecutionProvider"]) else: print("CPU") sess= ort.InferenceSession("model.onnx", None) import time start = time.time() output = sess.run( output_names=['labels', 'boxes', 'scores'], #output_names=None, input_feed={'images': im_data.data.numpy(), "orig_target_sizes": size.data.numpy()}, ) end = time.time() fps = 1.0 / (end - start) print(fps)
根据结果获取目标类别,标注框与得分,并将其绘制在每张图像上
labels, boxes, scores = output draw = ImageDraw.Draw(img)#生成ImageDraw对象,用于画图 thrh = 0.6 for i in range(im_data.shape[0]): scr = scores[i] lab = labels[i][scr > thrh] box = boxes[i][scr > thrh] for l, b in zip(lab, box): draw.rectangle(list(b), outline='red',)#ImageDraw对象画框并写入类别,因为原本输出的类别是0,1,2,需要使用定义的数组来获取其真实类别名称 draw.text((b[0], b[1] - 10), text=str(classes[l.item()]),font_size=16, fill='blue', ) images.append(img)#将画好的每张图加入到images数组中
视频合成,指定合成视频的名称,帧率等信息,将images中的图像合成为视频
video_name = 'output.mp4' fps = 30 # 每秒钟30帧 fourcc = cv2.VideoWriter_fourcc(*"mp4v") video = cv2.VideoWriter(video_name, fourcc, fps, (640,640)) # 合成视频 for i in range(len(images)): img = images[i] img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) video.write(img)
可以看到,其在视频推理过程中,对CPU的利用率明显增高
当然我们也可以直接将其生成gif
imageio.mimsave('output.gif',images,fps=25)
完整代码如下:
import torch import onnxruntime as ort from PIL import Image, ImageDraw from torchvision.transforms import ToTensor import cv2 import imageio import numpy as np import os if __name__ == "__main__": import torch print(torch.__version__) print(torch.version.cuda) print(torch.backends.cudnn.version()) ################## classes = ['car','truck',"bus"] ################## videoname = '1.mp4' capture = cv2.VideoCapture(videoname) images = [] if capture.isOpened(): while True: ret,img=capture.read() # img 就是一帧图片 if not ret:break # 当获取完最后一帧就结束 img=cv2.resize(img,(640,640)) img = Image.fromarray(img) img = img.convert('RGB') im_data = ToTensor()(img)[None] size = torch.tensor([[640, 640]]) if torch.cuda.is_available(): print("GPU") sess = ort.InferenceSession("model.onnx", None, providers=["CUDAExecutionProvider"]) else: print("CPU") sess= ort.InferenceSession("model.onnx", None) import time start = time.time() output = sess.run( output_names=['labels', 'boxes', 'scores'], #output_names=None, input_feed={'images': im_data.data.numpy(), "orig_target_sizes": size.data.numpy()}, ) end = time.time() fps = 1.0 / (end - start) print(fps) labels, boxes, scores = output draw = ImageDraw.Draw(img) thrh = 0.6 for i in range(im_data.shape[0]): scr = scores[i] lab = labels[i][scr > thrh] box = boxes[i][scr > thrh] for l, b in zip(lab, box): draw.rectangle(list(b), outline='red',) draw.text((b[0], b[1] - 10), text=str(classes[l.item()]),font_size=16, fill='blue', ) images.append(img) # 可以用 cv2.imshow() 查看这一帧,也可以逐帧保存 else: print('视频打开失败!') ############# img.save('2.jpg') imageio.mimsave('output.gif',images,fps=25) video_name = 'output.mp4' fps = 30 # 每秒钟30帧 fourcc = cv2.VideoWriter_fourcc(*"mp4v") video = cv2.VideoWriter(video_name, fourcc, fps, (640,640)) # 合成视频 for i in range(len(images)): img = images[i] img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) video.write(img)
文章版权声明:除非注明,否则均为主机测评原创文章,转载或复制请以超链接形式并注明出处。