def test_get_by_name_negative(self):
source = [
ClassItemMsg(id=1, name='one', color='red'),
ClassItemMsg(id=2, name='two', color='green'),
ClassItemMsg(id=3, name='three', color='blue')
]
cm = ClassMap.construct_from(source)
self.assertRaises(ValueError, lambda: cm.get_by_name('four'))
def test_all_color_false(self):
source = [
ClassItemMsg(id=1, name='one', color='red'),
ClassItemMsg(id=2, name='two', color='green'),
ClassItemMsg(id=3, name='three')
]
cm = ClassMap.construct_from(source)
self.assertFalse(cm.has_all_colors())
def test_category_index(self):
source = [
ClassItemMsg(id=1, name='one', color='red'),
ClassItemMsg(id=2, name='two', color='green'),
ClassItemMsg(id=3, name='three')
]
cm = ClassMap.construct_from(source)
index = cm.get_category_index()
self.assertEqual(index[1], {'id': 1, 'name': 'one'})
def test_construct_from_protos(self):
source = [
ClassItemMsg(id=1, name='one', color='red'),
ClassItemMsg(id=2, name='two', color='green'),
ClassItemMsg(id=3, name='three', color='blue')
]
cm = ClassMap.construct_from(source)
for i, msg in enumerate(source):
expected = ClassItem.from_proto(msg)
actual = cm.get_by_id(i + 1)
self.assertEqual(actual, expected)
def test_create_proto_from_task(self):
t = rv.TaskConfig.builder(rv.SEMANTIC_SEGMENTATION) \
.with_classes(['car', 'boat']) \
.with_chip_size(500) \
.build()
msg = t.to_proto()
expected_classes = [
ClassItemMsg(name='car', id=1),
ClassItemMsg(name='boat', id=2)
]
self.assertEqual(msg.task_type, rv.SEMANTIC_SEGMENTATION)
self.assertEqual(msg.semantic_segmentation_config.chip_size, 500)
actual_class_items = dict(
[(i.id, i) for i in msg.semantic_segmentation_config.class_items])
expected_class_items = dict([(i.id, i) for i in expected_classes])
self.assertDictEqual(actual_class_items, expected_class_items)
def to_proto(self):
return ClassItemMsg(id=self.id, name=self.name, color=self.color)