def test_keep_categories_with_unique_id(self): label_map_proto = string_int_label_map_pb2.StringIntLabelMap() label_map_string = """ item { id:2 name:'cat' } item { id:1 name:'child' } item { id:1 name:'person' } item { id:1 name:'n00007846' } """ text_format.Merge(label_map_string, label_map_proto) categories = label_map_util.convert_label_map_to_categories( label_map_proto, max_num_classes=3) self.assertListEqual([{ 'id': 2, 'name': u'cat' }, { 'id': 1, 'name': u'child' }], categories)
def test_convert_label_map_to_categories_with_few_classes(self): label_map_proto = self._generate_label_map(num_classes=4) cat_no_offset = label_map_util.convert_label_map_to_categories( label_map_proto, max_num_classes=2) expected_categories_list = [{ 'name': u'1', 'id': 1 }, { 'name': u'2', 'id': 2 }] self.assertListEqual(expected_categories_list, cat_no_offset)
def test_convert_label_map_to_categories_no_label_map(self): categories = label_map_util.convert_label_map_to_categories( None, max_num_classes=3) expected_categories_list = [{ 'name': u'category_1', 'id': 1 }, { 'name': u'category_2', 'id': 2 }, { 'name': u'category_3', 'id': 3 }] self.assertListEqual(expected_categories_list, categories)
# Path to label map file PATH_TO_LABELS = os.path.join(CWD_PATH, 'training', 'label.pbtxt') # Path to video PATH_TO_VIDEO = os.path.join(CWD_PATH, VIDEO_NAME) # Number of classes the object detector can identify NUM_CLASSES = 6 # Load the label map. # Label maps map indices to category names, so that when our convolution # network predicts `5`, we know that this corresponds to `king`. # Here we use internal utility functions, but anything that returns a # dictionary mapping integers to appropriate string labels would be fine label_map = label_map_util.load_labelmap(PATH_TO_LABELS) categories = label_map_util.convert_label_map_to_categories( label_map, max_num_classes=NUM_CLASSES, use_display_name=True) category_index = label_map_util.create_category_index(categories) # Load the Tensorflow model into memory. detection_graph = tf.Graph() with detection_graph.as_default(): od_graph_def = tf.GraphDef() with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid: serialized_graph = fid.read() od_graph_def.ParseFromString(serialized_graph) tf.import_graph_def(od_graph_def, name='') sess = tf.Session(graph=detection_graph) # Define input and output tensors (i.e. data) for the object detection classifier