如何使用tensorflow在视频上保存包含检测到的对象的帧?

时间:2018-01-21 08:05:05

标签: python opencv tensorflow object-detection

我使用tensorflow训练了自己的人物检测器,如何从视频中仅提取检测到的物体的帧? 我是python和tensorflow的新手,这是代码的最后一部分,我想检测发生了:

 import cv2
 cap = cv2.VideoCapture('http://192.168.15.8:8080/video')

 with detection_graph.as_default():
   with tf.Session(graph=detection_graph) as sess:
    ret = True
    while (ret):
       ret,image_np = cap.read()
       # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
       image_np_expanded = np.expand_dims(image_np, axis=0)
       image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
       # Each box represents a part of the image where a particular object was detected.
       boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
       # Each score represent how level of confidence for each of the objects.
       # Score is shown on the result image, together with the class label.
       scores = detection_graph.get_tensor_by_name('detection_scores:0')
       classes = detection_graph.get_tensor_by_name('detection_classes:0')
       num_detections = detection_graph.get_tensor_by_name('num_detections:0')
       # Actual detection.
       (boxes, scores, classes, num_detections) = sess.run(
           [boxes, scores, classes, num_detections],
           feed_dict={image_tensor: image_np_expanded})
       # Visualization of the results of a detection.
       vis_util.visualize_boxes_and_labels_on_image_array(
           image_np,
           np.squeeze(boxes),
           np.squeeze(classes).astype(np.int32),
           np.squeeze(scores),
           category_index,
           use_normalized_coordinates=True,
           line_thickness=8)

       cv2.imshow('image',cv2.resize(image_np,(600,400)))      

       if cv2.waitKey(25) & 0xFF == ord('q'):
           cv2.destroyAllWindows()
           cap.release()
           break

这是我的有点伪代码,假设上面代码中的“类”负责显示类的名称:

 count = 0
 if classes == true/1/person: #since my detector only has 1 class
    cv2.imwrite("frame%d.jpg" % count, image_np)              
    count = count + 1

 if vis_util.visualize_boxes_and_labels_on_image_array == true
     save frames

如何正确实施?

编辑:我已经通过在visualize_boxes_and_labels_on_image_array()函数中添加以下代码并在我的主脚本上添加计数器来解决我的问题。

if class_name == 'person':
  cv2.imwrite("frame%d.jpg" % count, image)

0 个答案:

没有答案