research/object_detection/utils/label_map_util_test.py
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for object_detection.utils.label_map_util."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import numpy as np
from six.moves import range
import tensorflow.compat.v1 as tf
from google.protobuf import text_format
from object_detection.protos import string_int_label_map_pb2
from object_detection.utils import label_map_util
class LabelMapUtilTest(tf.test.TestCase):
def _generate_label_map(self, num_classes):
label_map_proto = string_int_label_map_pb2.StringIntLabelMap()
for i in range(1, num_classes + 1):
item = label_map_proto.item.add()
item.id = i
item.name = 'label_' + str(i)
item.display_name = str(i)
return label_map_proto
def _generate_label_map_with_hierarchy(self, num_classes, ancestors_dict,
descendants_dict):
label_map_proto = string_int_label_map_pb2.StringIntLabelMap()
for i in range(1, num_classes + 1):
item = label_map_proto.item.add()
item.id = i
item.name = 'label_' + str(i)
item.display_name = str(i)
if i in ancestors_dict:
for anc_i in ancestors_dict[i]:
item.ancestor_ids.append(anc_i)
if i in descendants_dict:
for desc_i in descendants_dict[i]:
item.descendant_ids.append(desc_i)
return label_map_proto
def test_get_label_map_dict(self):
label_map_string = """
item {
id:2
name:'cat'
}
item {
id:1
name:'dog'
}
"""
label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
with tf.gfile.Open(label_map_path, 'wb') as f:
f.write(label_map_string)
label_map_dict = label_map_util.get_label_map_dict(label_map_path)
self.assertEqual(label_map_dict['dog'], 1)
self.assertEqual(label_map_dict['cat'], 2)
def test_get_keypoint_label_map_dict(self):
label_map_string = """
item: {
id: 1
name: 'face'
display_name: 'face'
keypoints {
id: 0
label: 'left_eye'
}
keypoints {
id: 1
label: 'right_eye'
}
}
item: {
id: 2
name: '/m/01g317'
display_name: 'person'
keypoints {
id: 2
label: 'left_shoulder'
}
keypoints {
id: 3
label: 'right_shoulder'
}
}
"""
label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
with tf.gfile.Open(label_map_path, 'wb') as f:
f.write(label_map_string)
label_map_dict = label_map_util.get_keypoint_label_map_dict(label_map_path)
self.assertEqual(label_map_dict['left_eye'], 0)
self.assertEqual(label_map_dict['right_eye'], 1)
self.assertEqual(label_map_dict['left_shoulder'], 2)
self.assertEqual(label_map_dict['right_shoulder'], 3)
def test_get_keypoint_label_map_dict_invalid(self):
label_map_string = """
item: {
id: 1
name: 'face'
display_name: 'face'
keypoints {
id: 0
label: 'left_eye'
}
keypoints {
id: 1
label: 'right_eye'
}
}
item: {
id: 2
name: '/m/01g317'
display_name: 'person'
keypoints {
id: 0
label: 'left_shoulder'
}
keypoints {
id: 1
label: 'right_shoulder'
}
}
"""
label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
with tf.gfile.Open(label_map_path, 'wb') as f:
f.write(label_map_string)
with self.assertRaises(ValueError):
_ = label_map_util.get_keypoint_label_map_dict(
label_map_path)
def test_get_label_map_dict_from_proto(self):
label_map_string = """
item {
id:2
name:'cat'
}
item {
id:1
name:'dog'
}
"""
label_map_proto = text_format.Parse(
label_map_string, string_int_label_map_pb2.StringIntLabelMap())
label_map_dict = label_map_util.get_label_map_dict(label_map_proto)
self.assertEqual(label_map_dict['dog'], 1)
self.assertEqual(label_map_dict['cat'], 2)
def test_get_label_map_dict_display(self):
label_map_string = """
item {
id:2
display_name:'cat'
}
item {
id:1
display_name:'dog'
}
"""
label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
with tf.gfile.Open(label_map_path, 'wb') as f:
f.write(label_map_string)
label_map_dict = label_map_util.get_label_map_dict(
label_map_path, use_display_name=True)
self.assertEqual(label_map_dict['dog'], 1)
self.assertEqual(label_map_dict['cat'], 2)
def test_load_bad_label_map(self):
label_map_string = """
item {
id:0
name:'class that should not be indexed at zero'
}
item {
id:2
name:'cat'
}
item {
id:1
name:'dog'
}
"""
label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
with tf.gfile.Open(label_map_path, 'wb') as f:
f.write(label_map_string)
with self.assertRaises(ValueError):
label_map_util.load_labelmap(label_map_path)
def test_load_label_map_with_background(self):
label_map_string = """
item {
id:0
name:'background'
}
item {
id:2
name:'cat'
}
item {
id:1
name:'dog'
}
"""
label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
with tf.gfile.Open(label_map_path, 'wb') as f:
f.write(label_map_string)
label_map_dict = label_map_util.get_label_map_dict(label_map_path)
self.assertEqual(label_map_dict['background'], 0)
self.assertEqual(label_map_dict['dog'], 1)
self.assertEqual(label_map_dict['cat'], 2)
def test_get_label_map_dict_with_fill_in_gaps_and_background(self):
label_map_string = """
item {
id:3
name:'cat'
}
item {
id:1
name:'dog'
}
"""
label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
with tf.gfile.Open(label_map_path, 'wb') as f:
f.write(label_map_string)
label_map_dict = label_map_util.get_label_map_dict(
label_map_path, fill_in_gaps_and_background=True)
self.assertEqual(label_map_dict['background'], 0)
self.assertEqual(label_map_dict['dog'], 1)
self.assertEqual(label_map_dict['2'], 2)
self.assertEqual(label_map_dict['cat'], 3)
self.assertEqual(len(label_map_dict), max(label_map_dict.values()) + 1)
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.Parse(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_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 test_convert_label_map_to_categories_lvis_frequency_and_counts(self):
label_map_proto = string_int_label_map_pb2.StringIntLabelMap()
label_map_string = """
item {
id:1
name:'person'
frequency: FREQUENT
instance_count: 1000
}
item {
id:2
name:'dog'
frequency: COMMON
instance_count: 100
}
item {
id:3
name:'cat'
frequency: RARE
instance_count: 10
}
"""
text_format.Parse(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': 1,
'name': u'person',
'frequency': 'f',
'instance_count': 1000
}, {
'id': 2,
'name': u'dog',
'frequency': 'c',
'instance_count': 100
}, {
'id': 3,
'name': u'cat',
'frequency': 'r',
'instance_count': 10
}], categories)
def test_convert_label_map_to_categories(self):
label_map_proto = self._generate_label_map(num_classes=4)
categories = label_map_util.convert_label_map_to_categories(
label_map_proto, max_num_classes=3)
expected_categories_list = [{
'name': u'1',
'id': 1,
}, {
'name': u'2',
'id': 2,
}, {
'name': u'3',
'id': 3,
}]
self.assertListEqual(expected_categories_list, categories)
def test_convert_label_map_with_keypoints_to_categories(self):
label_map_str = """
item {
id: 1
name: 'person'
keypoints: {
id: 1
label: 'nose'
}
keypoints: {
id: 2
label: 'ear'
}
}
"""
label_map_proto = string_int_label_map_pb2.StringIntLabelMap()
text_format.Parse(label_map_str, label_map_proto)
categories = label_map_util.convert_label_map_to_categories(
label_map_proto, max_num_classes=1)
self.assertEqual('person', categories[0]['name'])
self.assertEqual(1, categories[0]['id'])
self.assertEqual(1, categories[0]['keypoints']['nose'])
self.assertEqual(2, categories[0]['keypoints']['ear'])
def test_disallow_duplicate_keypoint_ids(self):
label_map_str = """
item {
id: 1
name: 'person'
keypoints: {
id: 1
label: 'right_elbow'
}
keypoints: {
id: 1
label: 'left_elbow'
}
}
item {
id: 2
name: 'face'
keypoints: {
id: 3
label: 'ear'
}
}
"""
label_map_proto = string_int_label_map_pb2.StringIntLabelMap()
text_format.Parse(label_map_str, label_map_proto)
with self.assertRaises(ValueError):
label_map_util.convert_label_map_to_categories(
label_map_proto, max_num_classes=2)
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_get_max_label_map_index(self):
num_classes = 4
label_map_proto = self._generate_label_map(num_classes=num_classes)
max_index = label_map_util.get_max_label_map_index(label_map_proto)
self.assertEqual(num_classes, max_index)
def test_create_category_index(self):
categories = [{'name': u'1', 'id': 1}, {'name': u'2', 'id': 2}]
category_index = label_map_util.create_category_index(categories)
self.assertDictEqual({
1: {
'name': u'1',
'id': 1
},
2: {
'name': u'2',
'id': 2
}
}, category_index)
def test_create_categories_from_labelmap(self):
label_map_string = """
item {
id:1
name:'dog'
}
item {
id:2
name:'cat'
}
"""
label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
with tf.gfile.Open(label_map_path, 'wb') as f:
f.write(label_map_string)
categories = label_map_util.create_categories_from_labelmap(label_map_path)
self.assertListEqual([{
'name': u'dog',
'id': 1
}, {
'name': u'cat',
'id': 2
}], categories)
def test_create_category_index_from_labelmap(self):
label_map_string = """
item {
id:2
name:'cat'
}
item {
id:1
name:'dog'
}
"""
label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
with tf.gfile.Open(label_map_path, 'wb') as f:
f.write(label_map_string)
category_index = label_map_util.create_category_index_from_labelmap(
label_map_path)
self.assertDictEqual({
1: {
'name': u'dog',
'id': 1
},
2: {
'name': u'cat',
'id': 2
}
}, category_index)
def test_create_category_index_from_labelmap_display(self):
label_map_string = """
item {
id:2
name:'cat'
display_name:'meow'
}
item {
id:1
name:'dog'
display_name:'woof'
}
"""
label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
with tf.gfile.Open(label_map_path, 'wb') as f:
f.write(label_map_string)
self.assertDictEqual({
1: {
'name': u'dog',
'id': 1
},
2: {
'name': u'cat',
'id': 2
}
}, label_map_util.create_category_index_from_labelmap(
label_map_path, False))
self.assertDictEqual({
1: {
'name': u'woof',
'id': 1
},
2: {
'name': u'meow',
'id': 2
}
}, label_map_util.create_category_index_from_labelmap(label_map_path))
def test_get_label_map_hierarchy_lut(self):
num_classes = 5
ancestors = {2: [1, 3], 5: [1]}
descendants = {1: [2], 5: [1, 2]}
label_map = self._generate_label_map_with_hierarchy(num_classes, ancestors,
descendants)
gt_hierarchy_dict_lut = {
'ancestors':
np.array([
[1, 0, 0, 0, 0],
[1, 1, 1, 0, 0],
[0, 0, 1, 0, 0],
[0, 0, 0, 1, 0],
[1, 0, 0, 0, 1],
]),
'descendants':
np.array([
[1, 1, 0, 0, 0],
[0, 1, 0, 0, 0],
[0, 0, 1, 0, 0],
[0, 0, 0, 1, 0],
[1, 1, 0, 0, 1],
]),
}
ancestors_lut, descendants_lut = (
label_map_util.get_label_map_hierarchy_lut(label_map, True))
np.testing.assert_array_equal(gt_hierarchy_dict_lut['ancestors'],
ancestors_lut)
np.testing.assert_array_equal(gt_hierarchy_dict_lut['descendants'],
descendants_lut)
if __name__ == '__main__':
tf.test.main()