Tensorflow Object Detection API 源码分析之 utils/label_map_util.py

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 !=