loading with empty annotations

luxonis · Nov 27, 2024 · b39761a · b39761a
1 parent abf2f9a
commit b39761a
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/luxonis_ml/data/loaders/luxonis_loader.py b/luxonis_ml/data/loaders/luxonis_loader.py
@@ -323,4 +323,21 @@ def _load_image_with_annotations(
 
             labels[task] = (array, anns[0]._label_type)
 
+        if not labels:
+            for task in self.classes_by_task.keys():
+                if task == LabelType.SEGMENTATION:
+                    empty_array = np.zeros(
+                        (len(self.class_mappings[task]), height, width),
+                        dtype=np.uint8,
+                    )
+                if task == LabelType.BOUNDINGBOX:
+                    empty_array = np.zeros((0, 6), dtype=np.float32)
+                if task == LabelType.KEYPOINTS:
+                    empty_array = np.zeros((0, 3), dtype=np.float32)
+                if task == LabelType.CLASSIFICATION:
+                    empty_array = np.zeros(
+                        (0, len(self.class_mappings[task])), dtype=np.float32
+                    )
+                labels[task] = (empty_array, task)
+
         return img, labels