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 main(unused_argv): assert FLAGS.checkpoint_dir, '`checkpoint_dir` is missing.' assert FLAGS.eval_dir, '`eval_dir` is missing.' tf.gfile.MakeDirs(FLAGS.eval_dir) if FLAGS.pipeline_config_path: configs = config_util.get_configs_from_pipeline_file( FLAGS.pipeline_config_path) tf.gfile.Copy(FLAGS.pipeline_config_path, os.path.join(FLAGS.eval_dir, 'pipeline.config'), overwrite=True) else: configs = config_util.get_configs_from_multiple_files( model_config_path=FLAGS.model_config_path, eval_config_path=FLAGS.eval_config_path, eval_input_config_path=FLAGS.input_config_path) for name, config in [('model.config', FLAGS.model_config_path), ('eval.config', FLAGS.eval_config_path), ('input.config', FLAGS.input_config_path)]: tf.gfile.Copy(config, os.path.join(FLAGS.eval_dir, name), overwrite=True) model_config = configs['model'] eval_config = configs['eval_config'] input_config = configs['eval_input_config'] if FLAGS.eval_training_data: input_config = configs['train_input_config'] model_fn = functools.partial(model_builder.build, model_config=model_config, is_training=False) def get_next(config): return dataset_builder.make_initializable_iterator( dataset_builder.build(config)).get_next() create_input_dict_fn = functools.partial(get_next, input_config) label_map = label_map_util.load_labelmap(input_config.label_map_path) max_num_classes = max([item.id for item in label_map.item]) categories = label_map_util.convert_label_map_to_categories( label_map, max_num_classes) if FLAGS.run_once: eval_config.max_evals = 1 graph_rewriter_fn = None if 'graph_rewriter_config' in configs: graph_rewriter_fn = graph_rewriter_builder.build( configs['graph_rewriter_config'], is_training=False) evaluator.evaluate(create_input_dict_fn, model_fn, eval_config, categories, FLAGS.checkpoint_dir, FLAGS.eval_dir, graph_hook_fn=graph_rewriter_fn)
def test_convert_label_map_to_coco_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 __fetch_category_indices(): dir_path = TFWeaponDetectionAPI.__get_dir_path() path_to_labels = os.path.join(dir_path + '/data', 'label_map.pbtxt') class_count = 1 label_map = label_map_util.load_labelmap(path_to_labels) categories = label_map_util.convert_label_map_to_categories( label_map, max_num_classes=class_count, use_display_name=True) category_index = label_map_util.create_category_index(categories) category_dict = {} for item in category_index.values(): category_dict[item['id']] = item['name'] category_dict[item['name']] = item['id'] return category_index, category_dict
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)
def __init__(self, model_name=PRETRAINED_ssd_mobilenet_v1_coco_2017_11_17): self.dir_path = dirname(realpath(__file__)) self.model_path = self.dir_path + '/object_detection/pretrained/' self.model_file = model_name + '.tar.gz' self.download_base = 'http://download.tensorflow.org/models/object_detection/' self.path_to_frozen_graph = self.model_path + model_name + '/frozen_inference_graph.pb' path_to_labels = os.path.join(self.dir_path + '/object_detection/data', 'mscoco_label_map.pbtxt') self.class_count = 90 if not path.exists(self.path_to_frozen_graph): self.__download() self.__load() self.label_map = label_map_util.load_labelmap(path_to_labels) self.categories = label_map_util.convert_label_map_to_categories( self.label_map, max_num_classes=self.class_count, use_display_name=True) self.category_index = label_map_util.create_category_index( self.categories) self.inPipe = Pipe() self.outPipe = Pipe()