11# Copyright (C) 2024 Intel Corporation
22# SPDX-License-Identifier: Apache-2.0
3+ from __future__ import annotations
34
4- from otx .core .types .label import NullLabelInfo , SegLabelInfo
5+ from datumaro import LabelCategories
6+ from datumaro .components .annotation import GroupType
7+ from otx .core .types .label import HLabelInfo , NullLabelInfo , SegLabelInfo
58
69
710def test_as_json (fxt_label_info ):
@@ -18,3 +21,34 @@ def test_seg_label_info():
1821 )
1922 assert SegLabelInfo .from_num_classes (1 ) == SegLabelInfo (["background" , "label_0" ], [["background" , "label_0" ]])
2023 assert SegLabelInfo .from_num_classes (0 ) == NullLabelInfo ()
24+
25+
26+ # Unit test
27+ def test_hlabel_info ():
28+ labels = [
29+ LabelCategories .Category (name = "car" , parent = "vehicle" ),
30+ LabelCategories .Category (name = "truck" , parent = "vehicle" ),
31+ LabelCategories .Category (name = "plush toy" , parent = "plush toy" ),
32+ LabelCategories .Category (name = "No class" ),
33+ ]
34+ label_groups = [
35+ LabelCategories .LabelGroup (
36+ name = "Detection labels___vehicle" ,
37+ labels = ["car" , "truck" ],
38+ group_type = GroupType .EXCLUSIVE ,
39+ ),
40+ LabelCategories .LabelGroup (
41+ name = "Detection labels___plush toy" ,
42+ labels = ["plush toy" ],
43+ group_type = GroupType .EXCLUSIVE ,
44+ ),
45+ LabelCategories .LabelGroup (name = "No class" , labels = ["No class" ], group_type = GroupType .RESTRICTED ),
46+ ]
47+ dm_label_categories = LabelCategories (items = labels , label_groups = label_groups )
48+
49+ hlabel_info = HLabelInfo .from_dm_label_groups (dm_label_categories )
50+
51+ # Check if class_to_group_idx and label_to_idx have the same keys
52+ assert list (hlabel_info .class_to_group_idx .keys ()) == list (
53+ hlabel_info .label_to_idx .keys (),
54+ ), "class_to_group_idx and label_to_idx keys do not match"
0 commit comments