Fix loading logic for datasets with multiple tasks

Author: JSabadin

luxonis_ml/data/loaders/luxonis_loader.py

@@ -323,11 +323,13 @@ def _load_image_with_annotations(
 
             labels[task] = (array, anns[0]._label_type)
 
-        if not labels:
-            for task in self.classes_by_task.keys():
+        missing_tasks = set(self.classes_by_task) - set(labels_by_task)
+        if missing_tasks:
+            for task in missing_tasks:
+                class_mapping_len = len(self.class_mappings[task])
                 if task == LabelType.SEGMENTATION:
                     empty_array = np.zeros(
-                        (len(self.class_mappings[task]), height, width),
+                        (class_mapping_len, height, width),
                         dtype=np.uint8,
                     )
                 elif task == LabelType.BOUNDINGBOX:
@@ -336,7 +338,7 @@ def _load_image_with_annotations(
                     empty_array = np.zeros((0, 3), dtype=np.float32)
                 elif task == LabelType.CLASSIFICATION:
                     empty_array = np.zeros(
-                        (0, len(self.class_mappings[task])), dtype=np.float32
+                        (0, class_mapping_len), dtype=np.float32
                     )
                 labels[task] = (empty_array, task)
 

tests/test_data/test_dataset.py

@@ -595,3 +595,63 @@ def generator():
     loader = LuxonisLoader(dataset, augmentations=augments)
     for _, labels in loader:
         assert labels == {}
+
+
+def test_partial_labels():
+    dataset = LuxonisDataset("__partial_labels", delete_existing=True)
+
+    def generator():
+        for i in range(8):
+            img = make_image(i)
+            if i < 2:
+                yield {
+                    "file": img,
+                }
+            elif i < 4:
+                yield {
+                    "file": img,
+                    "annotation": {
+                        "type": "classification",
+                        "class": "dog",
+                    },
+                }
+            elif i < 6:
+                yield {
+                    "file": img,
+                    "annotation": {
+                        "type": "boundingbox",
+                        "class": "dog",
+                        "x": 0.1,
+                        "y": 0.1,
+                        "w": 0.1,
+                        "h": 0.1,
+                    },
+                }
+                yield {
+                    "file": img,
+                    "annotation": {
+                        "type": "keypoints",
+                        "class": "dog",
+                        "keypoints": [[0.1, 0.1, 0], [0.2, 0.2, 1]],
+                    },
+                }
+            elif i < 8:
+                yield {
+                    "file": img,
+                    "annotation": {
+                        "type": "mask",
+                        "class": "dog",
+                        "mask": np.random.rand(512, 512) > 0.5,
+                    },
+                }
+
+    dataset.add(generator())
+    dataset.make_splits([1, 0, 0])
+
+    augments = Augmentations([512, 512], [{"name": "Rotate", "params": {}}])
+    loader = LuxonisLoader(dataset, augmentations=augments, view="train")
+    for _, labels in loader:
+        assert labels.get("boundingbox") is not None
+        assert labels.get("classification") is not None
+        assert labels.get("segmentation") is not None
+        assert labels.get("keypoints") is not None