最近接到一个需求,要分析视频中的物体,比如分析一段视频中是否有人,是否有车等等。
首先想到的是深度学习,机器学习,但是之前只是稍微看了看,没有深入学习,想要在短时间内搞定算法不太可能,于是就在github上搜索解决方案,找到不少,都是基于tensorflow的,比如yolo。
自己找了几个测试,发现这个东西太消耗性能,比如一段12s的1080p的视频,在2核4G加Tesla P40显卡的环境下,居然需要18s才能分析完成,计算这么一个视频,耗时18s,需要这么高的显卡,显然是不具备商业价值的,太贵。
后来找到一个分析图片的算法,自己改造成了分析视频的了,在4核8G无显卡的虚拟机中,仅需6s就可以分析完成如上视频,虽然识别率不如yolo高,但是目前看也够用了。如果再加入抽帧的算法,比如每隔10帧抽一帧进行分析,速度还能再提升几倍。现分享源代码如下,希望对各位同道中人有所帮助:
#!/usr/bin/env python # -*- coding:utf-8 -*- import numpy as np import cv2 #等待分析视频路径 video_path = "./demo.mp4" CLASSES = ["background", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"] COLORS = np.random.uniform(0, 255, size=(len(CLASSES), 3)) prototxt = 'demo.prototxt.txt' model = 'demo.caffemodel' video = video_path video_name = video_path.split('/')[-1] #分析结果视频路径 result_video = 'result/%s' %(video_name) #筛选,物品识别概率大于0.2的会话框,可以手动改这个数值 confidence_input = 0.2 net = cv2.dnn.readNetFromCaffe(prototxt, model) #读取视频 cap = cv2.VideoCapture(video) #获取视频fps fps_video = cap.get(cv2.CAP_PROP_FPS) #设置视频编码器 fourcc = cv2.VideoWriter_fourcc(*"DIVX") #设置视频写入参数 videoWriter = cv2.VideoWriter(result_video, fourcc, fps_video, (1920, 1080)) while (cap.isOpened()): ret, frame = cap.read() if ret == True: image = frame (h, w) = image.shape[:2] blob = cv2.dnn.blobFromImage(cv2.resize(image, (300, 300)), 0.007843, (300, 300), 127.5) net.setInput(blob) detections = net.forward() for i in np.arange(0, detections.shape[2]): confidence = detections[0, 0, i, 2] if confidence > confidence_input: idx = int(detections[0, 0, i, 1]) box = detections[0, 0, i, 3:7] * np.array([w, h, w, h]) (startX, startY, endX, endY) = box.astype("int") label = "{}: {:.2f}%".format(CLASSES[idx], confidence * 100) cv2.rectangle(image, (startX, startY), (endX, endY),COLORS[idx], 2) y = startY - 15 if startY - 15 > 15 else startY + 15 cv2.putText(image, label, (startX, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLORS[idx], 2) videoWriter.write(image) else: break videoWriter.release()
两个训练模型放在百度云
链接: https://pan.baidu.com/s/1Ozg3wgXMwlBeX4_joFVhKQ 密码: 75hw
需要的同学自取