我正在使用带有Open CV的TF Object Detection API

时间:2018-02-09 08:15:56

标签: python opencv tensorflow object-detection

如何提取视频检测到的对象类型。例如,一旦Object Detection API中的视频检测到“笔记本电脑”,我如何获得“笔记本”标签及其ID,以显示单独的文件?

import cv2
 cap = cv2.VideoCapture(0)

 with detection_graph.as_default():
   with tf.Session(graph=detection_graph) as sess:
    ret = True
    while (ret):
       ret,image_np = cap.read()

       image_np_expanded = np.expand_dims(image_np, axis=0)
       image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')

       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

1 个答案:

答案 0 :(得分:1)

假设您有标签映射的pbtxt文件,如下所示:

item {
  name: "/m/01g317"
  id: 1
  display_name: "person"
}
item {
  name: "/m/0199g"
  id: 2
  display_name: "bicycle"
}
item {
  name: "/m/0k4j"
  id: 3
  display_name: "car"
}
...

您可以使用 label_map_util [https://github.com/tensorflow/models/blob/master/research/object_detection/utils/label_map_util.py]

读取字典中的标签
label_map = label_map_util.load_labelmap(LABELS_PBTXT_FILE_PATH)

categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=n_classes, use_display_name=True)

idx_to_label = {}
for cat in categories:
    idx_to_label[cat['id']] = cat['name']

然后 - 当您拥有 idx_to_label 字典时,只需使用

idx_to_label.get(curr_id, 'N/A')