diff --git a/luxonis_ml/data/datasets/utils.py b/luxonis_ml/data/datasets/utils.py index e87cee8c..77483da1 100644 --- a/luxonis_ml/data/datasets/utils.py +++ b/luxonis_ml/data/datasets/utils.py @@ -239,7 +239,7 @@ def create_new_record( } for record in generator: - if isinstance(record, DatasetRecord): + if isinstance(record, DatasetRecord) or "annotation" not in record: yield record continue diff --git a/luxonis_ml/data/loaders/luxonis_loader.py b/luxonis_ml/data/loaders/luxonis_loader.py index 7fe086a8..951b8675 100644 --- a/luxonis_ml/data/loaders/luxonis_loader.py +++ b/luxonis_ml/data/loaders/luxonis_loader.py @@ -94,6 +94,8 @@ def __init__( self.class_mappings = {} for task in df["task"].unique(): + if not task: + continue class_mapping = { class_: i for i, class_ in enumerate( @@ -290,8 +292,9 @@ def _load_image_with_annotations( "instance_id": instance_id, } ) - annotation = load_annotation(type_, data) - labels_by_task[task].append(annotation) + if type_ != "NoneType": + annotation = load_annotation(type_, data) + labels_by_task[task].append(annotation) labels: Labels = {} for task, anns in labels_by_task.items():