Fixed records without annotations (#197)

luxonis · Oct 31, 2024 · e7121ff · e7121ff
1 parent 97d301a
commit e7121ff
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git 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
@@ -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():