viz_utils.visualize_boxes_and_labels_on_image_array get类别

Question

我是一个非常实用的土木工程师(非常糟糕的程序员)。我想检测建筑工地上的物体。现在我遵循教程，我得到了一张带有检测到的对象的边界框的图片。这里我的问题是:如何获取包含类别名称、索引或id的字符串？enter image description here

使用以下代码：

input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32)
detections = detect_fn(input_tensor)

num_detections = int(detections.pop('num_detections'))
detections = {key: value[0, :num_detections].numpy()
              for key, value in detections.items()}
detections['num_detections'] = num_detections

# detection_classes should be ints.
detections['detection_classes'] = detections['detection_classes'].astype(np.int64)

label_id_offset = 1
image_np_with_detections = image_np.copy()

viz_utils.visualize_boxes_and_labels_on_image_array(
            image_np_with_detections,
            detections['detection_boxes'],
            detections['detection_classes']+label_id_offset,
            detections['detection_scores'],
            category_index,
            use_normalized_coordinates=True,
            max_boxes_to_draw=5,
            min_score_thresh=.6,
            line_thickness=10,
            keypoint_edges=4,
            agnostic_mode=False)


plt.imshow(cv2.cvtColor(image_np_with_detections, cv2.COLOR_BGR2RGB))
plt.show()

Answer 1

您应该将category_index变量初始化为以下代码：

PATH_TO_LABELS = './models/research/object_detection/data/mscoco_label_map.pbtxt'
category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)
# category_index is a dictionary as seen below:
# {1: {"id": 1, "name": "category1"}, 2: {"id": 2, "name": "category2"}}

其中models目录指的是git clone --depth 1 https://github.com/tensorflow/models。

更多帮助可以在here上找到。