From b39761a7adcaaa552e27b6c4e555f1fd78b74527 Mon Sep 17 00:00:00 2001 From: Jernej Sabadin Date: Wed, 27 Nov 2024 06:48:04 +0100 Subject: [PATCH 1/4] loading with empty annotations --- luxonis_ml/data/loaders/luxonis_loader.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/luxonis_ml/data/loaders/luxonis_loader.py b/luxonis_ml/data/loaders/luxonis_loader.py index 8e0842d1..0c9b9e7e 100644 --- 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 From 3716ba708e96f8e4288e06e28f56bd19a4db5851 Mon Sep 17 00:00:00 2001 From: Jernej Sabadin Date: Wed, 27 Nov 2024 07:00:23 +0100 Subject: [PATCH 2/4] refactor conditional logic --- luxonis_ml/data/loaders/luxonis_loader.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/luxonis_ml/data/loaders/luxonis_loader.py b/luxonis_ml/data/loaders/luxonis_loader.py index 0c9b9e7e..44853a9c 100644 --- a/luxonis_ml/data/loaders/luxonis_loader.py +++ b/luxonis_ml/data/loaders/luxonis_loader.py @@ -330,11 +330,11 @@ def _load_image_with_annotations( (len(self.class_mappings[task]), height, width), dtype=np.uint8, ) - if task == LabelType.BOUNDINGBOX: + elif task == LabelType.BOUNDINGBOX: empty_array = np.zeros((0, 6), dtype=np.float32) - if task == LabelType.KEYPOINTS: + elif task == LabelType.KEYPOINTS: empty_array = np.zeros((0, 3), dtype=np.float32) - if task == LabelType.CLASSIFICATION: + elif task == LabelType.CLASSIFICATION: empty_array = np.zeros( (0, len(self.class_mappings[task])), dtype=np.float32 ) From 62e6cb9e94149622a1c5d7110d641a2be2771ce6 Mon Sep 17 00:00:00 2001 From: Jernej Sabadin Date: Thu, 28 Nov 2024 05:55:56 +0100 Subject: [PATCH 3/4] Fix loading logic for datasets with multiple tasks --- luxonis_ml/data/loaders/luxonis_loader.py | 10 ++-- tests/test_data/test_dataset.py | 60 +++++++++++++++++++++++ 2 files changed, 66 insertions(+), 4 deletions(-) diff --git a/luxonis_ml/data/loaders/luxonis_loader.py b/luxonis_ml/data/loaders/luxonis_loader.py index 44853a9c..77779ef2 100644 --- a/luxonis_ml/data/loaders/luxonis_loader.py +++ b/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) diff --git a/tests/test_data/test_dataset.py b/tests/test_data/test_dataset.py index 210c1aec..f733769b 100644 --- a/tests/test_data/test_dataset.py +++ b/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 From 1fe1a3d599cb50cfe10ee8995f9994a72e937d79 Mon Sep 17 00:00:00 2001 From: Jernej Sabadin Date: Thu, 28 Nov 2024 06:51:55 +0100 Subject: [PATCH 4/4] fix failing tests --- luxonis_ml/data/loaders/luxonis_loader.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/luxonis_ml/data/loaders/luxonis_loader.py b/luxonis_ml/data/loaders/luxonis_loader.py index 77779ef2..3c14eb4d 100644 --- a/luxonis_ml/data/loaders/luxonis_loader.py +++ b/luxonis_ml/data/loaders/luxonis_loader.py @@ -324,7 +324,14 @@ 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: + 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: