深度学习分析识别视频中的物体

Posted 2021-01-02 zhenglisai

最近接到一个需求，要分析视频中的物体，比如分析一段视频中是否有人，是否有车等等。

首先想到的是深度学习，机器学习，但是之前只是稍微看了看，没有深入学习，想要在短时间内搞定算法不太可能，于是就在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

需要的同学自取