Handle dict style names in load_labelmap

Author: ford-downer-robo

roboflow/util/image_utils.py

@@ -102,8 +102,11 @@ def load_labelmap(f):
     if f.lower().endswith(".yaml") or f.lower().endswith(".yml"):
         with open(f) as file:
             data = yaml.safe_load(file)
-            if "names" in data:
-                return {i: name for i, name in enumerate(data["names"])}
+            names = data.get("names", [])
+            if isinstance(names, dict):
+                return {int(k): v for k, v in names.items()}
+            else:
+                return {i: name for i, name in enumerate(names)}
     else:
         with open(f) as file:
             lines = [line for line in file.readlines() if line.strip()]

tests/annotations/dict_names.yaml

@@ -0,0 +1,5 @@
+names:
+  0: cat
+  1: dog
+  2: fish
+nc: 3

tests/util/test_image_utils.py

@@ -2,7 +2,7 @@
 
 import responses
 
-from roboflow.util.image_utils import check_image_path, check_image_url
+from roboflow.util.image_utils import check_image_path, check_image_url, load_labelmap
 
 
 class TestCheckImagePath(unittest.TestCase):
@@ -33,3 +33,9 @@ def test_url_not_found(self):
         url = "https://example.com/notfound.png"
         responses.add(responses.HEAD, url, status=404)
         self.assertFalse(check_image_url(url))
+
+
+class TestLoadLabelmap(unittest.TestCase):
+    def test_yaml_dict_names(self):
+        labelmap = load_labelmap("tests/annotations/dict_names.yaml")
+        self.assertEqual(labelmap, {0: "cat", 1: "dog", 2: "fish"})