深度学习实战（十四）——Tensorflow Object Detection API 源码分析之 utils/label_map_util.py

以你之姓@ 2024-04-19 09:44 17阅读 0赞

## Tensorflow Object Detection API 源码分析之 utils/label\_map\_util.py ##

`protos/string_int_label_map.proto`

syntax = "proto2";
    
    package object_detection.protos;
    
    message StringIntLabelMapItem {
      // String name. The most common practice is to set this to a MID or synsets
      // id.
      optional string name = 1;
    
      // Integer id that maps to the string name above. Label ids should start from
      // 1.
      optional int32 id = 2;
    
      // Human readable string label.
      optional string display_name = 3;
    };
    
    message StringIntLabelMap {
      repeated StringIntLabelMapItem item = 1;
    };

utils/label\_map\_util.py

# 有关 label map 的辅助函数
    """Label map utility functions."""
    
    import logging
    
    import tensorflow as tf
    from google.protobuf import text_format
    from object_detection.protos import string_int_label_map_pb2
    
    
    # background默认为0，所以其他的不能为0，否则报错
    def _validate_label_map(label_map):
      """Checks if a label map is valid.
    
      Args:
        label_map: StringIntLabelMap to validate.
    
      Raises:
        ValueError: if label map is invalid.
      """
      for item in label_map.item:
        if item.id < 0:
          raise ValueError('Label map ids should be >= 0.')
        if (item.id == 0 and item.name != 'background' and
            item.display_name != 'background'):
          raise ValueError('Label map id 0 is reserved for the background label')
    
    
    # 返回 类别 和 index 的字典  {key为index，value为类别id}
    # 例如 category_index[1] = id ?
    def create_category_index(categories):
      """Creates dictionary of COCO compatible categories keyed by category id.
    
      Args:
        categories: a list of dicts, each of which has the following keys:
          'id': (required) an integer id uniquely identifying this category.
          'name': (required) string representing category name
            e.g., 'cat', 'dog', 'pizza'.
    
      Returns:
        category_index: a dict containing the same entries as categories, but keyed
          by the 'id' field of each category.
      """
      category_index = {}
      for cat in categories:
        category_index[cat['id']] = cat
      return category_index
    
    # 返回最大的index
    def get_max_label_map_index(label_map):
      """Get maximum index in label map.
    
      Args:
        label_map: a StringIntLabelMapProto
    
      Returns:
        an integer
      """
      return max([item.id for item in label_map.item])
    
    
    # 将label_map 转化为 categories 列表 [{'id': item.id, 'name': name} , ...]
    # 可选 name 或者 display_name
    def convert_label_map_to_categories(label_map,
                                        max_num_classes,
                                        use_display_name=True):
      """Loads label map proto and returns categories list compatible with eval.
    
      This function loads a label map and returns a list of dicts, each of which
      has the following keys:
        'id': (required) an integer id uniquely identifying this category.
        'name': (required) string representing category name
          e.g., 'cat', 'dog', 'pizza'.
      We only allow class into the list if its id-label_id_offset is
      between 0 (inclusive) and max_num_classes (exclusive).
      If there are several items mapping to the same id in the label map,
      we will only keep the first one in the categories list.
    
      Args:
        label_map: a StringIntLabelMapProto or None.  If None, a default categories
          list is created with max_num_classes categories.
        max_num_classes: maximum number of (consecutive) label indices to include.
        use_display_name: (boolean) choose whether to load 'display_name' field
          as category name.  If False or if the display_name field does not exist,
          uses 'name' field as category names instead.
      Returns:
        categories: a list of dictionaries representing all possible categories.
      """
      categories = []
      list_of_ids_already_added = []
      if not label_map:
        label_id_offset = 1
        for class_id in range(max_num_classes):
          categories.append({
              'id': class_id + label_id_offset,
              'name': 'category_{}'.format(class_id + label_id_offset)
          })
        return categories
      for item in label_map.item:
        if not 0 < item.id <= max_num_classes:
          logging.info('Ignore item %d since it falls outside of requested '
                       'label range.', item.id)
          continue
        if use_display_name and item.HasField('display_name'):
          name = item.display_name
        else:
          name = item.name
        if item.id not in list_of_ids_already_added:
          list_of_ids_already_added.append(item.id)
          categories.append({'id': item.id, 'name': name})
      return categories
    
    
    # 从文件载入 label map，下一个函数使用
    def load_labelmap(path):
      """Loads label map proto.
    
      Args:
        path: path to StringIntLabelMap proto text file.
      Returns:
        a StringIntLabelMapProto
      """
      with tf.gfile.GFile(path, 'r') as fid:
        label_map_string = fid.read()
        label_map = string_int_label_map_pb2.StringIntLabelMap()
        try:
          text_format.Merge(label_map_string, label_map)
        except text_format.ParseError:
          label_map.ParseFromString(label_map_string)
      _validate_label_map(label_map)
      return label_map
    
    
    # 获取 label_map 字典，返回形式 label_map_dict['background'] = 0
    def get_label_map_dict(label_map_path,
                           use_display_name=False,
                           fill_in_gaps_and_background=False):
      """Reads a label map and returns a dictionary of label names to id.
    
      Args:
        label_map_path: path to StringIntLabelMap proto text file.
        use_display_name: whether to use the label map items' display names as keys.
        fill_in_gaps_and_background: whether to fill in gaps and background with
        respect to the id field in the proto. The id: 0 is reserved for the
        'background' class and will be added if it is missing. All other missing
        ids in range(1, max(id)) will be added with a dummy class name
        ("class_<id>") if they are missing.
    
      Returns:
        A dictionary mapping label names to id.
    
      Raises:
        ValueError: if fill_in_gaps_and_background and label_map has non-integer or
        negative values.
      """
      label_map = load_labelmap(label_map_path)
      label_map_dict = {}
      for item in label_map.item:
        if use_display_name:
          label_map_dict[item.display_name] = item.id
        else:
          label_map_dict[item.name] = item.id
    
      if fill_in_gaps_and_background:
        values = set(label_map_dict.values())
    
        if 0 not in values:
          label_map_dict['background'] = 0
        if not all(isinstance(value, int) for value in values):
          raise ValueError('The values in label map must be integers in order to'
                           'fill_in_gaps_and_background.')
        if not all(value >= 0 for value in values):
          raise ValueError('The values in the label map must be positive.')
    
        if len(values) != max(values) + 1:
          # there are gaps in the labels, fill in gaps.
          for value in range(1, max(values)):
            if value not in values:
              label_map_dict['class_' + str(value)] = value
    
      return label_map_dict
    
    
    # 建立从类别索引 到 label_map的字典  
    # {1: {'id': 1, 'name': 'dog'}, 2: {'id': 2, 'name': 'cat'}, ...}
    # model_lib.py 中调用的函数，重要
    def create_category_index_from_labelmap(label_map_path):
      """Reads a label map and returns a category index.
    
      Args:
        label_map_path: Path to `StringIntLabelMap` proto text file.
    
      Returns:
        A category index, which is a dictionary that maps integer ids to dicts
        containing categories, e.g.
        {1: {'id': 1, 'name': 'dog'}, 2: {'id': 2, 'name': 'cat'}, ...}
      """
      label_map = load_labelmap(label_map_path)
      max_num_classes = max(item.id for item in label_map.item)
      categories = convert_label_map_to_categories(label_map, max_num_classes)
      return create_category_index(categories)
    
    # 建立类别无关的索引，只有一个object物体的类别
    def create_class_agnostic_category_index():
      """Creates a category index with a single `object` class."""
      return {1: {'id': 1, 'name': 'object'}}

转自：

[Tensorflow Object Detection API 源码分析之 utils/label\_map\_util.py][Tensorflow Object Detection API _ utils_label_map_util.py]：这个人的博客里还有关于其它目标检测API的代码

[Tensorflow Object Detection API _ utils_label_map_util.py]: https://blog.csdn.net/lifengcai_/article/details/81703945