Fix augmented loader with no labels

luxonis_ml/data/loaders/luxonis_loader.py

@@ -184,6 +184,11 @@
         loaded_anns = [self._load_image_with_annotations(i) for i in indices]
         random_state = random.getstate()
         np_random_state = np.random.get_state()
+        if not loaded_anns[0][1]:
+            img, aug_annotations = self.augmentations(
+                [(loaded_anns[i][0], {}) for i in range(len(loaded_anns))],
+            )
+
         while loaded_anns[0][1]:
             aug_input_data = []
             label_to_task = {}
@@ -242,7 +247,7 @@
             for label_type, array in aug_annotations.items():
                 out_dict[label_to_task[label_type]] = (array, label_type)
 
-        return img, out_dict  # type: ignore
+        return img, out_dict
 
     def _load_image_with_annotations(
         self, idx: int
@@ -256,7 +261,13 @@
             with all the present annotations
         """
 
+        if not self.idx_to_df_row:
+            raise ValueError(
+                f"No data found in dataset '{self.dataset.identifier}' for {self.view} views"
+            )
+
         ann_indices = self.idx_to_df_row[idx]
+
         ann_rows = [self.df.row(row) for row in ann_indices]
         if not self.dataset.is_remote:
             img_path = ann_rows[0][8]

tests/test_data/test_dataset.py

@@ -572,3 +572,26 @@ def generator():
         assert "array" in labels
         assert labels["label"][0].tolist() == [["dog"], ["cat"]]
         assert np.allclose(labels["array"][0], arr)
+
+
+def test_no_labels():
+    dataset = LuxonisDataset("__no_labels", delete_existing=True)
+
+    def generator():
+        for i in range(10):
+            img = make_image(i)
+            yield {
+                "file": img,
+            }
+
+    dataset.add(generator())
+    dataset.make_splits()
+
+    loader = LuxonisLoader(dataset)
+    for _, labels in loader:
+        assert labels == {}
+
+    augments = Augmentations([512, 512], [{"name": "Flip", "params": {}}])
+    loader = LuxonisLoader(dataset, augmentations=augments)
+    for _, labels in loader:
+        assert labels == {}