opencv读取视频使用ssd处理定位类别并将处理后类别存储起来

使用需要的环境

• Tensorflow
• opencv3

tensorflow 中的ssd目标定位检测算法下载链接
文件下载好后将下面代码中的
file_dir 改成你下载的SSD-Tensorflow文件目录，
path 改成你需要处理的文件目录
src 为你要输出的目录
更改完成后就可以运行程序

import os
import math
import random

import numpy as np
import tensorflow as tf
import cv2

slim = tf.contrib.slim

import matplotlib.image as mpimg

import sys

# 你需要更改的内容
#SSD-Tensorflow-master文件夹下载的主目录位置
file_dir = "/home/rui/Files/SSD-Tensorflow-master/"
#你需要识别的文件目录
path = '/home/rui/Downloads/新警察故事BD国粤双语中字[电影天堂www.dy2018.com].mkv'
#你的输出文件名称
src = "output.avi"
###

sys.path.append(file_dir)

from nets import ssd_vgg_300, ssd_common, np_methods
from preprocessing import ssd_vgg_preprocessing

# TensorFlow session: grow memory when needed. TF, DO NOT USE ALL MY GPU MEMORY!!!
gpu_options = tf.GPUOptions(allow_growth=True)
config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options)
isess = tf.InteractiveSession(config=config)

# Input placeholder.
net_shape = (300, 300)
data_format = 'NHWC'
img_input = tf.placeholder(tf.uint8, shape=(None, None, 3))
# Evaluation pre-processing: resize to SSD net shape.
image_pre, labels_pre, bboxes_pre, bbox_img = ssd_vgg_preprocessing.preprocess_for_eval(
    img_input, None, None, net_shape, data_format, resize=ssd_vgg_preprocessing.Resize.WARP_RESIZE)
image_4d = tf.expand_dims(image_pre, 0)

# Define the SSD model.
reuse = True if 'ssd_net' in locals() else None
ssd_net = ssd_vgg_300.SSDNet()
with slim.arg_scope(ssd_net.arg_scope(data_format=data_format)):
    predictions, localisations, _, _ = ssd_net.net(image_4d, is_training=False, reuse=reuse)
# Restore SSD model.
ckpt_filename = file_dir + "checkpoints/ssd_300_vgg.ckpt"
# ckpt_filename = '../checkpoints/VGG_VOC0712_SSD_300x300_ft_iter_120000.ckpt'
isess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(isess, ckpt_filename)

# SSD default anchor boxes.
ssd_anchors = ssd_net.anchors(net_shape)


# Main image processing routine.
def process_image(img, select_threshold=0.5, nms_threshold=.45, net_shape=(300, 300)):
    # Run SSD network.
    rimg, rpredictions, rlocalisations, rbbox_img = isess.run([image_4d, predictions, localisations, bbox_img],
                                                              feed_dict={img_input: img})

    # Get classes and bboxes from the net outputs.
    rclasses, rscores, rbboxes = np_methods.ssd_bboxes_select(
        rpredictions, rlocalisations, ssd_anchors,
        select_threshold=select_threshold, img_shape=net_shape, num_classes=21, decode=True)

    rbboxes = np_methods.bboxes_clip(rbbox_img, rbboxes)
    rclasses, rscores, rbboxes = np_methods.bboxes_sort(rclasses, rscores, rbboxes, top_k=400)
    rclasses, rscores, rbboxes = np_methods.bboxes_nms(rclasses, rscores, rbboxes, nms_threshold=nms_threshold)
    # Resize bboxes to original image shape. Note: useless for Resize.WARP!
    rbboxes = np_methods.bboxes_resize(rbbox_img, rbboxes)
    return rclasses, rscores, rbboxes



def plt_bboxes(img, classes, scores, bboxes, thickness=0):
    """
    plt_bboxes(img, classes, scores, bboxes)
    .   @brief Draws a simple, thick, or filled up-right rectangle and filled in classes and scores.
    .
    .
    .   @param img Image.
    .   @param classes Vertex of the rectangle.
    .   @param scores Vertex of the rectangle opposite to pt1 .
    .   @param bboxes Rectangle color or brightness (grayscale image).
    .   @param thickness Thickness of lines that make up the rectangle. Negative values, like #FILLED,
    """

    labels = [["Null" , (30, 180, 0)],
              ["Aeroplanes", (255, 0, 0)],
              ["Bicycles", (150, 0, 0)],
              ["Birds", (70, 0, 0)],
              ["Boats", (0, 0, 0)],
              ["Bottles", (0, 70, 0)],
              ["Buses", (0, 150, 0)],
              ["Cars", (0, 255, 0)],
              ["Cats", (0, 0, 70)],
              ["Chairs", (0, 0, 150)],
              ["Cows", (0, 55, 55)],
              ["Dining tables", (55, 55, 55)],
              ["Dogs", (55, 0, 255)],
              ["Horses", (100, 100, 0)],
              ["Motorbikes", (0, 100, 100)],
              ["People", (0, 0, 255)],
              ["Potted plants", (100, 0, 100)],
              ["Sheep", (0, 150, 150)],
              ["Sofas", (150, 0, 150)],
              ["Trains",(150, 150, 0)],
              ["TV/Monitors",(200, 200, 0)]]
    img_shape = img.shape
    cols = img_shape[0]
    rows = img_shape[1]


    length = len(classes)
    for i in range(length):
        label = labels[classes[i]][0]
        color = labels[classes[i]][1]
        lecolor = labels[classes[i] - 1][1]
        score = scores[i]
        Letter = label + "|" + str(score)
#        Letter = bboxes[][] - bboxes[][]
        p1 = (int(bboxes[i][1] * rows), int(bboxes[i][0] * cols))
        p2 = (int(bboxes[i][3] * rows), int(bboxes[i][2] * cols))
        cv2.rectangle(img, p1, p2, color, thickness=thickness)
        font = cv2.FONT_HERSHEY_DUPLEX
        cv2.putText(img, Letter, p1, font, 0.5, lecolor, 1, cv2.LINE_AA)





def process_video(src, dst, is_show=False):
    videoCapture = cv2.VideoCapture(src)
    fps = videoCapture.get(cv2.CAP_PROP_FPS)
    length = int(videoCapture.get(cv2.CAP_PROP_FRAME_COUNT))
    size = (int(videoCapture.get(cv2.CAP_PROP_FRAME_WIDTH)),
            int(videoCapture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    #fourcc = cv2.CAP_PROP_FOURCC()
    fourcc = cv2.VideoWriter_fourcc(*'X264')
    videoWriter = cv2.VideoWriter(dst, fourcc, fps, size)

    ret, frame = videoCapture.read()
    sum = 0
    while ret:
        rclasses, rscores, rbboxes = process_image(frame)
        plt_bboxes(frame, rclasses, rscores, rbboxes)

        if is_show:
            cv2.imshow("windows", frame)
            cv2.waitKey(1)
        videoWriter.write(frame)  # 写视频帧
        success, frame = videoCapture.read()
        sum = sum + 1
#        print('' % sum*100/length)
        pro = sum*100/length
        print(('processing %.2f{}' % pro).format("%"))
        length
    videoCapture.release()
    videoWriter.release()

    # Test on some demo image and visualize output.

process_video(path, src)


# visualization.bboxes_draw_on_img(img, rclasses, rscores, rbboxes, visualization.colors_plasma)