diff --git a/luxonis_ml/data/loaders/luxonis_loader.py b/luxonis_ml/data/loaders/luxonis_loader.py index 951b8675..1711fd4f 100644 --- a/luxonis_ml/data/loaders/luxonis_loader.py +++ b/luxonis_ml/data/loaders/luxonis_loader.py @@ -184,6 +184,11 @@ def __getitem__(self, idx: int) -> LuxonisLoaderOutput: 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 @@ def __getitem__(self, idx: int) -> LuxonisLoaderOutput: 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 @@ def _load_image_with_annotations( 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] diff --git a/tests/test_data/test_dataset.py b/tests/test_data/test_dataset.py index 3935ecef..210c1aec 100644 --- a/tests/test_data/test_dataset.py +++ b/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 == {}