loading with empty annotations

luxonis_ml/data/loaders/luxonis_loader.py

@@ -323,4 +323,30 @@ def _load_image_with_annotations(
 
             labels[task] = (array, anns[0]._label_type)
 
+        missing_tasks = set(self.classes_by_task) - set(labels_by_task)
+        if missing_tasks and missing_tasks.issubset(
+            {
+                LabelType.SEGMENTATION,
+                LabelType.BOUNDINGBOX,
+                LabelType.KEYPOINTS,
+                LabelType.CLASSIFICATION,
+            }
+        ):
+            for task in missing_tasks:
+                class_mapping_len = len(self.class_mappings[task])
+                if task == LabelType.SEGMENTATION:
+                    empty_array = np.zeros(
+                        (class_mapping_len, height, width),
+                        dtype=np.uint8,
+                    )
+                elif task == LabelType.BOUNDINGBOX:
+                    empty_array = np.zeros((0, 6), dtype=np.float32)
+                elif task == LabelType.KEYPOINTS:
+                    empty_array = np.zeros((0, 3), dtype=np.float32)
+                elif task == LabelType.CLASSIFICATION:
+                    empty_array = np.zeros(
+                        (0, class_mapping_len), dtype=np.float32
+                    )
+                labels[task] = (empty_array, task)
+
         return img, labels

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