Repository: luxonis/luxonis-ml (language: python)

Issue:
Version 0.1.0
# Version `0.1.0`

- NN Archive
- Test Coverage
  - code coverage 9 -> 66 %
- `LuxonisParser`
  - automatic creation of LDF from standard dataset formats
- `BaseDataset` and dataset registry for custom implementations
- New CLI


Discussion / hints:
## Test Results
53 tests   53 :white_check_mark:  47s :stopwatch:
 1 suites   0 :zzz:
 1 files     0 :x:

Results for commit eb5c61d3.

[test-results]:data:application/gzip;base64,H4sIAHELEGYC/12Myw6DIBQFf8Ww7gKEi9KfaRAhualKw2Nl+u9F21Lt7sycZFbicLKRXBt2aUjMmCqMOeiEfikousLlSdsH/Au3mI35M3d8FEOrcBqnk7Ah+PAxIS+1uO1T8C1+vZ0PuZ2PNePnGVMBYgcwko28U0ZQzoFB63Sves2UE+DEALSVVCryfAFrwlpJAAEAAA==

:recycle: This comment has been updated with latest results.
# ☂️ Python Coverage
> current status: ✅
## Overall Coverage
| Lines | Covered | Coverage | Threshold | Status |
| :---: | :-----: | :------: | :-------: | :----: |
|  3358 |   2216  |    66%   |     0%    |   🟢   |
## New Files
| File                                                                          | Coverage | Status |
| :---------------------------------------------------------------------------- | :------: | :----: |
| luxonis_ml/data/datasets/__init__.py                                          |   100%   |   🟢   |
| luxonis_ml/data/datasets/base_dataset.py                                      |    76%   |   🟢   |
| luxonis_ml/data/datasets/source.py                                            |    86%   |   🟢   |
| luxonis_ml/data/parsers/__init__.py                                           |   100%   |   🟢   |
| luxonis_ml/data/parsers/base_parser.py                                        |    90%   |   🟢   |
| luxonis_ml/data/parsers/classification_directory_parser.py                    |    97%   |   🟢   |
| luxonis_ml/data/parsers/coco_parser.py                                        |    80%   |   🟢   |
| luxonis_ml/data/parsers/create_ml_parser.py                                   |    97%   |   🟢   |
| luxonis_ml/data/parsers/darknet_parser.py                                     |    98%   |   🟢   |
| luxonis_ml/data/parsers/luxonis_parser.py                                     |    86%   |   🟢   |
| luxonis_ml/data/parsers/segmentation_mask_directory_parser.py                 |    96%   |   🟢   |
| luxonis_ml/data/parsers/solo_parser.py                                        |    19%   |   🟢   |
| luxonis_ml/data/parsers/tensorflow_csv_parser.py                              |   100%   |   🟢   |
| luxonis_ml/data/parsers/voc_parser.py                                         |    97%   |   🟢   |
| luxonis_ml/data/parsers/yolov4_parser.py                                      |   100%   |   🟢   |
| luxonis_ml/data/parsers/yolov6_parser.py                                      |    94%   |   🟢   |
| luxonis_ml/nn_archive/__init__.py                                             |   100%   |   🟢   |
| luxonis_ml/nn_archive/archive_generator.py                                    |    97%   |   🟢   |
| luxonis_ml/nn_archive/config.py                                               |   100%   |   🟢   |
| luxonis_ml/nn_archive/config_building_blocks/__init__.py                      |   100%   |   🟢   |
| luxonis_ml/nn_archive/config_building_blocks/base_models/__init__.py          |   100%   |   🟢   |
| luxonis_ml/nn_archive/config_building_blocks/base_models/custom_base_model.py |   100%   |   🟢   |
| luxonis_ml/nn_archive/config_building_blocks/base_models/head.py              |    77%   |   🟢   |
| luxonis_ml/nn_archive/config_building_blocks/base_models/head_outputs.py      |   100%   |   🟢   |
| luxonis_ml/nn_archive/config_building_blocks/base_models/input.py             |   100%   |   🟢   |
| luxonis_ml/nn_archive/config_building_blocks/base_models/metadata.py          |   100%   |   🟢   |
| luxonis_ml/nn_archive/config_building_blocks/base_models/output.py            |   100%   |   🟢   |
| luxonis_ml/nn_archive/config_building_blocks/enums/__init__.py                |   100%   |   🟢   |
| luxonis_ml/nn_archive/config_building_blocks/enums/data_type.py               |   100%   |   🟢   |
| luxonis_ml/nn_archive/config_building_blocks/enums/decoding.py                |   100%   |   🟢   |
| luxonis_ml/nn_archive/config_building_blocks/enums/input_type.py              |   100%   |   🟢   |
| luxonis_ml/nn_archive/model.py                                                |   100%   |   🟢   |
| luxonis_ml/nn_archive/utils.py                                                |    92%   |   🟢   |
| **TOTAL**                                                                     |  **93%** |   🟢   |
## Modified Files
| File                                           | Coverage | Status |
| :--------------------------------------------- | :------: | :----: |
| luxonis_ml/data/__init__.py                    |    86%   |   🟢   |
| luxonis_ml/data/augmentations/custom/mosaic.py |    69%   |   🟢   |
| luxonis_ml/data/loader.py                      |    82%   |   🟢   |
| luxonis_ml/data/utils/parquet.py               |    90%   |   🟢   |
| luxonis_ml/enums/__init__.py                   |   100%   |   🟢   |
| luxonis_ml/enums/enums.py                      |   100%   |   🟢   |
| luxonis_ml/guard_extras.py                     |    75%   |   🟢   |
| luxonis_ml/utils/__init__.py                   |   100%   |   🟢   |
| luxonis_ml/utils/environ.py                    |    97%   |   🟢   |
| luxonis_ml/utils/filesystem.py                 |    75%   |   🟢   |
| **TOTAL**                                      |  **87%** |   🟢   |


> **updated for commit: `eb5c61d` by [action](https://github.com/marketplace/actions/python-coverage)🐍**

Fix patch:
diff --git a/examples/Data_Custom_Example.ipynb b/examples/Data_Custom_Example.ipynb
index 006b106a..592b87bd 100644
--- a/examples/Data_Custom_Example.ipynb
+++ b/examples/Data_Custom_Example.ipynb
@@ -223,7 +223,7 @@
     "        }\n",
     "    }\n",
     ")\n",
-    "dataset.add(COCO_people_subset_generator)"
+    "dataset.add(COCO_people_subset_generator())"
    ]
   },
   {
@@ -262,6 +262,8 @@
    "source": [
     "loader = LuxonisLoader(dataset, view=\"train\")\n",
     "for image, ann in loader:\n",
+    "    # selecting the \"default\" task group\n",
+    "    ann = ann[\"default\"]\n",
     "    cls = ann[LabelType.CLASSIFICATION]\n",
     "    box = ann[LabelType.BOUNDINGBOX]\n",
     "    seg = ann[LabelType.SEGMENTATION]\n",
@@ -308,7 +310,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.8.19"
   }
  },
  "nbformat": 4,
diff --git a/examples/Data_Parser_Example.ipynb b/examples/Data_Parser_Example.ipynb
index 85d8634c..45f01e4b 100644
--- a/examples/Data_Parser_Example.ipynb
+++ b/examples/Data_Parser_Example.ipynb
@@ -110,6 +110,7 @@
    "source": [
     "loader = LuxonisLoader(dataset, view=\"train\")\n",
     "for image, ann in loader:\n",
+    "    ann = ann[\"default\"]\n",
     "    cls = ann[LabelType.CLASSIFICATION]\n",
     "    box = ann[LabelType.BOUNDINGBOX]\n",
     "    seg = ann[LabelType.SEGMENTATION]\n",
@@ -156,7 +157,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.18"
+   "version": "3.8.19"
   }
  },
  "nbformat": 4,
diff --git a/luxonis_ml/data/__init__.py b/luxonis_ml/data/__init__.py
index 2020b34a..e1dad68b 100644
--- a/luxonis_ml/data/__init__.py
+++ b/luxonis_ml/data/__init__.py
@@ -7,8 +7,7 @@
     from .datasets import (
         DATASETS_REGISTRY,
         BaseDataset,
-        DatasetGenerator,
-        DatasetGeneratorFunction,
+        DatasetIterator,
         LuxonisComponent,
         LuxonisDataset,
         LuxonisSource,
@@ -38,8 +37,7 @@ def load_dataset_plugins() -> None:
     "BaseLoader",
     "BucketStorage",
     "BucketType",
-    "DatasetGenerator",
-    "DatasetGeneratorFunction",
+    "DatasetIterator",
     "DATASETS_REGISTRY",
     "ImageType",
     "LabelType",
diff --git a/luxonis_ml/data/__main__.py b/luxonis_ml/data/__main__.py
index 2a81fb0d..f51ae847 100644
--- a/luxonis_ml/data/__main__.py
+++ b/luxonis_ml/data/__main__.py
@@ -39,27 +39,29 @@ def check_exists(name: str):
         raise typer.Exit()
 
 
-def get_dataset_info(name: str) -> Tuple[int, List[str], List[str]]:
+def get_dataset_info(name: str) -> Tuple[int, List[str], List[str], List[str]]:
     dataset = LuxonisDataset(name)
     size = len(dataset)
     try:
         loader = LuxonisLoader(dataset, view=SplitType.TRAIN.value)
         _, ann = next(iter(loader))
     except Exception:
-        ann = []
+        ann = {}
     classes, _ = dataset.get_classes()
-    tasks = [task.value for task in ann]
-    return size, classes, tasks
+    groups = list(ann.keys())
+    tasks = [task.value for group in groups for task in ann[group]]
+    return size, classes, tasks, groups
 
 
 def print_info(name: str) -> None:
-    size, classes, tasks = get_dataset_info(name)
+    size, classes, tasks, groups = get_dataset_info(name)
     print(
         Panel.fit(
             f"[magenta b]Name: [not b cyan]{name}\n"
             f"[magenta b]Size: [not b cyan]{size}\n"
             f"[magenta b]Classes: [not b cyan]{', '.join(classes)}\n"
-            f"[magenta b]Tasks: [not b cyan]{', '.join(tasks)}",
+            f"[magenta b]Tasks: [not b cyan]{', '.join(tasks)}\n"
+            f"[magenta b]Groups: [not b cyan]{', '.join(groups)}",
             title="Dataset Info",
         )
     )
@@ -104,7 +106,7 @@ def ls(
         dataset = LuxonisDataset(name)
         rows = [name, str(len(dataset))]
         if full:
-            _, classes, tasks = get_dataset_info(name)
+            _, classes, tasks, _ = get_dataset_info(name)
             rows.extend(
                 [
                     ", ".join(classes) if classes else "[red]<empty>[no red]",
@@ -134,7 +136,8 @@ def inspect(
     dataset = LuxonisDataset(name)
     loader = LuxonisLoader(dataset, view=view.value)
     for image, ann in loader:
-        # cls = ann[LabelType.CLASSIFICATION]
+        if len(ann) > 1:
+            raise NotImplementedError("Only one annotation group is supported.")
 
         h, w, _ = image.shape
         if LabelType.BOUNDINGBOX in ann:
diff --git a/luxonis_ml/data/augmentations/utils.py b/luxonis_ml/data/augmentations/utils.py
index 76480173..c7073ed4 100644
--- a/luxonis_ml/data/augmentations/utils.py
+++ b/luxonis_ml/data/augmentations/utils.py
@@ -105,7 +105,7 @@ def _parse_cfg(
 
     def __call__(
         self,
-        data: List[Tuple[np.ndarray, Dict[LabelType, np.ndarray]]],
+        data: List[Tuple[np.ndarray, Dict[str, Dict[LabelType, np.ndarray]]]],
         nc: int = 1,
         ns: int = 1,
         nk: int = 1,
diff --git a/luxonis_ml/data/datasets/__init__.py b/luxonis_ml/data/datasets/__init__.py
index 861453f1..4c7aea1a 100644
--- a/luxonis_ml/data/datasets/__init__.py
+++ b/luxonis_ml/data/datasets/__init__.py
@@ -1,16 +1,14 @@
 from .base_dataset import (
     DATASETS_REGISTRY,
     BaseDataset,
-    DatasetGenerator,
-    DatasetGeneratorFunction,
+    DatasetIterator,
 )
 from .luxonis_dataset import LuxonisDataset
 from .source import LuxonisComponent, LuxonisSource
 
 __all__ = [
     "BaseDataset",
-    "DatasetGenerator",
-    "DatasetGeneratorFunction",
+    "DatasetIterator",
     "LuxonisDataset",
     "LuxonisComponent",
     "LuxonisSource",
diff --git a/luxonis_ml/data/datasets/base_dataset.py b/luxonis_ml/data/datasets/base_dataset.py
index 2a432f74..238b9f6f 100644
--- a/luxonis_ml/data/datasets/base_dataset.py
+++ b/luxonis_ml/data/datasets/base_dataset.py
@@ -1,17 +1,139 @@
+import json
+import os.path as osp
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, Generator, List, Optional, Tuple
-
+from datetime import datetime
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, cast
+
+from pydantic import BaseModel, ConfigDict, model_validator
+from typeguard import TypeCheckError, check_type
+from typing_extensions import TypeAlias
+
+from luxonis_ml.data.utils.data_utils import (
+    ArrayType,
+    BoxType,
+    ClassificationType,
+    KeypointsType,
+    LabelType,
+    PolylineType,
+    SegmentationType,
+    transform_segmentation_value,
+)
+from luxonis_ml.enums import AnnotationType
 from luxonis_ml.utils.registry import AutoRegisterMeta, Registry
 
 from .source import LuxonisSource
 
 DATASETS_REGISTRY = Registry(name="datasets")
 
-DatasetGenerator = Generator[Dict[str, Any], None, None]
-DatasetGeneratorFunction = Callable[[], DatasetGenerator]
 
+DatasetIterator: TypeAlias = Iterator[dict]
+
+
+class Annotation(BaseModel):
+    """Base class for annotations in a dataset."""
+
+    model_config: ConfigDict = ConfigDict(extra="forbid")
+
+    file: str
+    class_: str
+    type_: AnnotationType
+    value: Union[
+        ClassificationType,
+        LabelType,
+        BoxType,
+        SegmentationType,
+        PolylineType,
+        KeypointsType,
+        ArrayType,
+    ]
+    task_group: str = "default"
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Annotation":
+        """Creates an annotation from a dictionary.
 
-class BaseDataset(ABC, metaclass=AutoRegisterMeta, registry=DATASETS_REGISTRY):
+        @type data: Dict[str, Any]
+        @param data: A dictionary of annotation data.
+        @rtype: L{Annotation}
+        @return: An annotation object.
+        """
+        return cls(
+            **{(f"{k}_" if k in ["class", "type"] else k): v for k, v in data.items()}
+        )
+
+    def to_parquet(self, instance_id: str) -> Dict[str, Union[str, datetime, None]]:
+        """Converts an annotation to a dictionary for writing to a parquet file.
+
+        @type instance_id: str
+        @param instance_id: Unique identifier for the instance
+        @rtype: Dict[str, str]
+        @return: A dictionary of annotation data.
+        """
+        value = self.value
+
+        if self.type_ == AnnotationType.SEGMENTATION:
+            value = transform_segmentation_value(cast(SegmentationType, self.value))
+
+        if isinstance(value, (list, tuple)):
+            value = json.dumps(value)
+        else:
+            value = str(value)
+
+        return {
+            "instance_id": instance_id,
+            "file": osp.basename(self.file),
+            "value_type": type(self.value).__name__,
+            "type": self.type_.value,
+            "class": self.class_,
+            "value": value,
+            "task_group": self.task_group,
+            "created_at": datetime.utcnow(),
+        }
+
+    @model_validator(mode="after")
+    def _value(self):
+        self.check_value_type(self.type_, self.value)
+        return self
+
+    @staticmethod
+    def check_value_type(typ: AnnotationType, value: Any) -> None:
+        def _check_value_type(expected_type: Any, value: Any) -> None:
+            """Checks if a value is of a given type, and raises a TypeError if not."""
+            try:
+                check_type(value, expected_type)
+            except TypeCheckError as e:
+                raise TypeError(
+                    f"Value {value} for '{typ.value}' is not of type {expected_type}"
+                ) from e
+
+        for expected_type, annotation_type in zip(
+            [
+                ClassificationType,
+                LabelType,
+                BoxType,
+                SegmentationType,
+                PolylineType,
+                KeypointsType,
+                ArrayType,
+            ],
+            [
+                AnnotationType.CLASSIFICATION,
+                AnnotationType.LABEL,
+                AnnotationType.BOX,
+                AnnotationType.SEGMENTATION,
+                AnnotationType.POLYLINE,
+                AnnotationType.KEYPOINTS,
+                AnnotationType.ARRAY,
+            ],
+        ):
+            if typ == annotation_type:
+                _check_value_type(expected_type, value)
+                break
+
+
+class BaseDataset(
+    ABC, metaclass=AutoRegisterMeta, registry=DATASETS_REGISTRY, register=False
+):
     """Base abstract dataset class for managing datasets in the Luxonis MLOps
     ecosystem."""
 
@@ -83,10 +205,10 @@ def get_skeletons(self) -> Dict[str, Dict]:
         pass
 
     @abstractmethod
-    def add(self, generator: DatasetGeneratorFunction) -> None:
+    def add(self, generator: DatasetIterator) -> None:
         """Write annotations to parquet files.
 
-        @type generator: L{DatasetGeneratorFunction}
+        @type generator: L{DatasetGenerator}
         @param generator: A Python iterator that yields dictionaries of data
             with the key described by the C{ANNOTATIONS_SCHEMA} but also listed below:
                 - file (C{str}) : path to file on local disk or object storage
diff --git a/luxonis_ml/data/datasets/luxonis_dataset.py b/luxonis_ml/data/datasets/luxonis_dataset.py
index 57edf0ad..7a2f2156 100644
--- a/luxonis_ml/data/datasets/luxonis_dataset.py
+++ b/luxonis_ml/data/datasets/luxonis_dataset.py
@@ -4,7 +4,7 @@
 import os.path as osp
 import shutil
 import time
-from datetime import datetime
+from contextlib import contextmanager
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
@@ -15,12 +15,13 @@
 import rich.progress
 
 import luxonis_ml.data.utils.data_utils as data_utils
+from luxonis_ml.enums import AnnotationType
 from luxonis_ml.utils import LuxonisFileSystem, environ
 
 from ..utils.constants import LABEL_TYPES, LDF_VERSION
 from ..utils.enums import BucketStorage, BucketType, ImageType, MediaType
 from ..utils.parquet import ParquetFileManager
-from .base_dataset import BaseDataset, DatasetGeneratorFunction
+from .base_dataset import Annotation, BaseDataset, DatasetIterator
 from .source import LuxonisComponent, LuxonisSource
 
 
@@ -281,12 +282,11 @@ def _write_index(
         table = pa.Table.from_pandas(df)
         pq.write_table(table, file_index_path)
 
-    def _start_time(self) -> None:
-        self.t0 = time.time()
-
-    def _end_time(self) -> None:
-        self.t1 = time.time()
-        self.logger.info(f"Took {self.t1 - self.t0} seconds")
+    @contextmanager
+    def _log_time(self):
+        t = time.time()
+        yield
+        self.logger.info(f"Took {time.time() - t} seconds")
 
     def _make_temp_dir(self) -> None:
         if osp.exists(self.tmp_dir):
@@ -429,12 +429,12 @@ def delete_dataset(self) -> None:
 
     def add(
         self,
-        generator: DatasetGeneratorFunction,
+        generator: DatasetIterator,
         batch_size: int = 1000000,
     ) -> None:
         """Write annotations to parquet files.
 
-        @type generator: L{DatasetGeneratorFunction}
+        @type generator: L{DatasetGenerator}
         @param generator: A Python iterator that yields dictionaries of data
             with the key described by the C{ANNOTATIONS_SCHEMA} but also listed below:
                 - file (C{str}) : path to file on local disk or object storage
@@ -459,25 +459,28 @@ def add(
             This can be set to a lower value to reduce memory usage.
         """
 
-        def _add_process_batch(batch_data: List[Dict]) -> None:
-            paths = list(set([data["file"] for data in batch_data]))
+        def _add_process_batch(batch_data: List[Annotation]) -> None:
+            paths = list(set(data.file for data in batch_data))
             self.logger.info("Generating UUIDs...")
-            self._start_time()
-            uuid_dict = self.fs.get_file_uuids(
-                paths, local=True
-            )  # TODO: support from bucket
-            self._end_time()
+            with self._log_time():
+                uuid_dict = self.fs.get_file_uuids(
+                    paths, local=True
+                )  # TODO: support from bucket
             if self.bucket_storage != BucketStorage.LOCAL:
                 self.logger.info("Uploading media...")
                 # TODO: support from bucket (likely with a self.fs.copy_dir)
-                self._start_time()
-                self.fs.put_dir(
-                    local_paths=paths, remote_dir="media", uuid_dict=uuid_dict
-                )
-                self._end_time()
+
+                with self._log_time():
+                    self.fs.put_dir(
+                        local_paths=paths, remote_dir="media", uuid_dict=uuid_dict
+                    )
 
             array_paths = list(
-                set([data["value"] for data in batch_data if data["type"] == "array"])
+                set(
+                    str(data.value)
+                    for data in batch_data
+                    if data.type_ == AnnotationType.ARRAY
+                )
             )
             progress = rich.progress.Progress(
                 rich.progress.TextColumn("[progress.description]{task.description}"),
@@ -489,78 +492,61 @@ def _add_process_batch(batch_data: List[Dict]) -> None:
             task = progress.add_task(
                 "[magenta]Processing data...", total=len(batch_data)
             )
-            if len(array_paths):
+            if array_paths:
                 self.logger.info("Checking arrays...")
-                self._start_time()
-                data_utils.check_arrays(array_paths)
-                self._end_time()
+                with self._log_time():
+                    data_utils.check_arrays(array_paths)
                 self.logger.info("Generating array UUIDs...")
-                self._start_time()
-                array_uuid_dict = self.fs.get_file_uuids(
-                    array_paths, local=True
-                )  # TODO: support from bucket
-                self._end_time()
+                with self._log_time():
+                    array_uuid_dict = self.fs.get_file_uuids(
+                        array_paths, local=True
+                    )  # TODO: support from bucket
                 if self.bucket_storage != BucketStorage.LOCAL:
                     self.logger.info("Uploading arrays...")
                     # TODO: support from bucket (likely with a self.fs.copy_dir)
-                    self._start_time()
-                    mask_upload_dict = self.fs.put_dir(
-                        local_paths=array_paths,
-                        remote_dir="arrays",
-                        uuid_dict=array_uuid_dict,
-                    )
-                    self._end_time()
+                    with self._log_time():
+                        mask_upload_dict = self.fs.put_dir(
+                            local_paths=array_paths,
+                            remote_dir="arrays",
+                            uuid_dict=array_uuid_dict,
+                        )
                 self.logger.info("Finalizing paths...")
                 progress.start()
-                for data in batch_data:
-                    if data["type"] == "array":
+                for ann in batch_data:
+                    if ann.type_ == AnnotationType.ARRAY:
                         if self.bucket_storage != BucketStorage.LOCAL:
-                            remote_path = mask_upload_dict[data["value"]]
+                            remote_path = mask_upload_dict[str(ann.value)]
                             remote_path = f"{self.fs.protocol}://{osp.join(self.fs.path, remote_path)}"
-                            data["value"] = remote_path
+                            ann.value = remote_path
                         else:
-                            data["value"] = osp.abspath(data["value"])
+                            ann.value = osp.abspath(str(ann.value))
                         progress.update(task, advance=1)
                 progress.stop()
 
             self.logger.info("Saving annotations...")
-            self._start_time()
-            progress.reset(task)
-            progress.start()
-            for data in batch_data:
-                filepath = data["file"]
-                file = osp.basename(filepath)
-                instance_id = uuid_dict[filepath]
-                # check for duplicate instance_ids to get a one-to-one relationship
-                matched_id = self._find_filepath_instance_id(filepath, index)
-                if matched_id is not None:
-                    if matched_id == instance_id:
-                        raise Exception(
-                            f"{filepath} already added to the dataset! Please skip or rename the file."
-                        )
-                elif instance_id not in new_index["instance_id"]:
-                    new_index["instance_id"].append(instance_id)
-                    new_index["file"].append(file)
-                    new_index["original_filepath"].append(osp.abspath(filepath))
-
-                data_utils.check_annotation(data)
-                data["instance_id"] = instance_id
-                data["file"] = file
-                data["value_type"] = type(data["value"]).__name__
-                if data["type"] == "segmentation":  # handles RLE
-                    data["value"] = data_utils.transform_segmentation_value(
-                        data["value"]
-                    )
-                if isinstance(data["value"], (list, tuple)):
-                    data["value"] = json.dumps(data["value"])  # convert lists to string
-                else:
-                    data["value"] = str(data["value"])
-                data["created_at"] = datetime.utcnow()
-
-                self.pfm.write(data)
-                progress.update(task, advance=1)
-            progress.stop()
-            self._end_time()
+            with self._log_time():
+                progress.reset(task)
+                progress.start()
+                for ann in batch_data:
+                    filepath = ann.file
+                    file = osp.basename(filepath)
+                    instance_id = uuid_dict[filepath]
+                    matched_id = self._find_filepath_instance_id(filepath, index)
+                    if matched_id is not None:
+                        if matched_id != instance_id:
+                            # TODO: not sure if this should be an exception or how we should really handle it
+                            raise Exception(
+                                f"{filepath} already added to the dataset! Please skip or rename the file."
+                            )
+                            # TODO: we may also want to check for duplicate instance_ids to get a one-to-one relationship
+                    elif instance_id not in new_index["instance_id"]:
+                        new_index["instance_id"].append(instance_id)
+                        new_index["file"].append(file)
+                        new_index["original_filepath"].append(osp.abspath(filepath))
+
+                    self.pfm.write(ann.to_parquet(instance_id))
+                    progress.update(task, advance=1)
+                progress.stop()
 
         if self.bucket_storage == BucketStorage.LOCAL:
             self.pfm = ParquetFileManager(self.annotations_path)
@@ -573,10 +559,10 @@ def _add_process_batch(batch_data: List[Dict]) -> None:
         index = self._get_file_index()
         new_index = {"instance_id": [], "file": [], "original_filepath": []}
 
-        batch_data = []
+        batch_data: list[Annotation] = []
 
-        for i, data in enumerate(generator()):
-            batch_data.append(data)
+        for i, data in enumerate(generator):
+            batch_data.append(Annotation.from_dict(data))
             if (i + 1) % batch_size == 0:
                 _add_process_batch(batch_data)
                 batch_data = []
@@ -629,6 +615,7 @@ def make_splits(
                 self.fs.get_dir("annotations", osp.join(self.tmp_dir, "annotations"))
 
             df = self._load_df_offline()
+            assert df is not None
             ids = list(set(df["instance_id"]))
             np.random.shuffle(ids)
             N = len(ids)
diff --git a/luxonis_ml/data/datasets/source.py b/luxonis_ml/data/datasets/source.py
index d7f6c57d..36504d54 100644
--- a/luxonis_ml/data/datasets/source.py
+++ b/luxonis_ml/data/datasets/source.py
@@ -1,40 +1,26 @@
+from dataclasses import dataclass
 from typing import List, Optional
 
 from ..utils.enums import ImageType, MediaType
 
 
+@dataclass
 class LuxonisComponent:
-    def __init__(
-        self,
-        name: str,
-        media_type: MediaType = MediaType.IMAGE,
-        image_type: Optional[ImageType] = ImageType.COLOR,
-    ) -> None:
-        """Abstraction for a piece of media within a source. Most commonly, this
-        abstracts an image sensor.
-
-        @type name: str
-        @param name: A recognizable name for the component.
-        @type media_type: MediaType
-        @param media_type: Enum for the type of media for the component.
-        @type image_type: Optional[ImageType]
-        @param image_type: Enum for the image type if C{media_type==MediaType.IMAGE}.
-            Else this param can be None or default.
-        """
-
-        if media_type not in MediaType:
-            raise Exception(f"{media_type.value} is not a valid MediaType")
-        if image_type not in ImageType:
-            raise Exception(f"{image_type.value} is not a valid ImageType")
-
-        self.name = name
-        self.media_type = media_type
-        self.image_type = image_type
-
-        if media_type == MediaType.IMAGE:
-            self.image_type = image_type
-        else:
-            self.image_type = None
+    """Abstraction for a piece of media within a source. Most commonly, this abstracts
+    an image sensor.
+
+    @type name: str
+    @param name: A recognizable name for the component.
+    @type media_type: L{MediaType}
+    @param media_type: Enum for the type of media for the component.
+    @type image_type: L{ImageType}
+    @param image_type: Enum for the image type. Only used if
+        C{media_type==MediaType.IMAGE}.
+    """
+
+    name: str
+    media_type: MediaType = MediaType.IMAGE
+    image_type: ImageType = ImageType.COLOR
 
 
 class LuxonisSource:
diff --git a/luxonis_ml/data/loader.py b/luxonis_ml/data/loader.py
index 95538e93..6ef98253 100644
--- a/luxonis_ml/data/loader.py
+++ b/luxonis_ml/data/loader.py
@@ -5,13 +5,14 @@
 import random
 import warnings
 from abc import ABC, abstractmethod
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 import cv2
 import numpy as np
 import pandas as pd
 import pycocotools.mask as mask_util
 from PIL import Image, ImageDraw
+from typing_extensions import TypeAlias
 
 from luxonis_ml.enums import LabelType
 
@@ -19,12 +20,54 @@
 from .datasets import LuxonisDataset
 from .utils.enums import BucketStorage
 
-Labels = Dict[LabelType, np.ndarray]
+Labels: TypeAlias = Dict[LabelType, np.ndarray]
 """C{Labels} is a dictionary of a label type and its annotations as L{numpy
 arrays<np.ndarray>}."""
 
-LuxonisLoaderOutput = Tuple[np.ndarray, Labels]
-"""C{LuxonisLoaderOutput} is a tuple of image and its annotations."""
+
+class Annotations(dict):
+    """A dictionary of task group names and their annotations.
+
+    Inherits from the built-in Python dictionary class.
+
+    Acts like simple dictionary when only one group is present, otherwise acts like a
+    nested dictionary of task group names and their annotations.
+    """
+
+    @property
+    def groups(self) -> List[str]:
+        return list(self.keys())
+
+    def __getitem__(self, key: Union[LabelType, str]) -> Union[np.ndarray, Labels]:
+        if isinstance(key, LabelType):
+            if len(self.groups) == 1:
+                return super().__getitem__(self.groups[0])[key]
+            raise ValueError(
+                "Multiple groups present, please access the specific group first."
+            )
+        elif isinstance(key, str):
+            if key in self.groups:
+                return super().__getitem__(key)
+            raise KeyError(f"Group '{key}' not found in annotations.")
+        else:
+            raise TypeError("Key must be of type LabelType or str.")
+
+    def __contains__(self, key: Union[LabelType, str]) -> bool:
+        if isinstance(key, LabelType):
+            if len(self.groups) == 1:
+                return key in super().__getitem__(self.groups[0])
+            raise ValueError(
+                "Multiple groups present, please access the specific group first."
+            )
+        elif isinstance(key, str):
+            return key in self.groups
+        else:
+            raise TypeError("Key must be of type LabelType or str.")
+
+
+LuxonisLoaderOutput: TypeAlias = Tuple[np.ndarray, Annotations]
+"""C{LuxonisLoaderOutput} is a tuple of an image as a L{numpy array<np.ndarray>} and a
+dictionary of task group names and their annotations as L{Annotations}."""
 
 
 class BaseLoader(ABC):
@@ -69,10 +112,10 @@ def __init__(
         @type view: str
         @param view: View of the dataset. Defaults to "train".
         @type stream: bool
-        @param stream: Flag for data streaming. Defaults to False.
+        @param stream: Flag for data streaming. Defaults to C{False}.
         @type augmentations: Optional[luxonis_ml.loader.Augmentations]
         @param augmentations: Augmentation class that performs augmentations. Defaults
-            to None.
+            to C{None}.
         """
 
         self.logger = logging.getLogger(__name__)
@@ -93,9 +136,10 @@ def __init__(
                 self.dataset.set_classes(classes=synced_classes[task], task=task)
 
         if self.dataset.bucket_storage == BucketStorage.LOCAL or not self.stream:
-            self.file_index = self.dataset._get_file_index()
-            if self.file_index is None:
+            file_index = self.dataset._get_file_index()
+            if file_index is None:
                 raise Exception("Cannot find file index")
+            self.file_index = file_index
         else:
             raise NotImplementedError(
                 "Streaming for remote bucket storage not implemented yet"
@@ -130,7 +174,10 @@ def __init__(
         else:
             raise NotImplementedError
 
-        self.df = dataset._load_df_offline(sync_mode=self.sync_mode)
+        df = dataset._load_df_offline(sync_mode=self.sync_mode)
+        if df is None:
+            raise Exception("Cannot find dataframe")
+        self.df = df
         self.df.set_index(["instance_id"], inplace=True)
 
     def __len__(self) -> int:
@@ -152,33 +199,104 @@ def __getitem__(self, idx: int) -> LuxonisLoaderOutput:
             annotations.
         """
 
-        img, annotations = self._load_image_with_annotations(idx)
-
-        if self.augmentations is not None:
-            aug_input_data = [(img, annotations)]
-            if self.augmentations.is_batched:
-                other_indices = [i for i in range(len(self)) if i != idx]
-                if self.augmentations.aug_batch_size > len(self):
-                    warnings.warn(
-                        f"Augmentations batch_size ({self.augmentations.aug_batch_size}) is larger than dataset size ({len(self)}), samples will include repetitions."
+        img, group_annotations = self._load_image_with_annotations(idx)
+
+        for task_group in list(group_annotations.keys()):
+            annotations = group_annotations[task_group]
+            if self.augmentations is not None:
+                aug_input_data = [(img, annotations)]
+                if self.augmentations.is_batched:
+                    other_indices = [i for i in range(len(self)) if i != idx]
+                    if self.augmentations.aug_batch_size > len(self):
+                        warnings.warn(
+                            f"Augmentations batch_size ({self.augmentations.aug_batch_size}) is larger than dataset size ({len(self)}), samples will include repetitions."
+                        )
+                        random_fun = random.choices
+                    else:
+                        random_fun = random.sample
+                    picked_indices = random_fun(
+                        other_indices, k=self.augmentations.aug_batch_size - 1
                     )
-                    random_fun = random.choices
-                else:
-                    random_fun = random.sample
-                picked_indices = random_fun(
-                    other_indices, k=self.augmentations.aug_batch_size - 1
-                )
-                aug_input_data.extend(
-                    [self._load_image_with_annotations(i) for i in picked_indices]
-                )
+                    for i in picked_indices:
+                        _img, _group_annotations = self._load_image_with_annotations(i)
+                        aug_input_data.append((_img, _group_annotations[task_group]))
 
-            img, annotations = self.augmentations(
-                aug_input_data, nc=self.nc, ns=self.ns, nk=self.max_nk
+                img, annotations = self.augmentations(
+                    aug_input_data, nc=self.nc, ns=self.ns, nk=self.max_nk
+                )
+                group_annotations[task_group] = annotations
+
+        return img, Annotations(group_annotations)
+
+    def _load_classification(self, rows: pd.DataFrame) -> np.ndarray:
+        classes = [row[1]["class"] for row in rows.iterrows() if bool(row[1]["value"])]
+        classify = np.zeros(self.nc)
+        for class_ in classes:
+            class_ = self.classes_by_task[LabelType.CLASSIFICATION].index(class_)
+            classify[class_] = classify[class_] + 1
+        classify[classify > 0] = 1
+        return classify
+
+    def _load_box(self, rows: pd.DataFrame) -> np.ndarray:
+        boxes = np.zeros((0, 5))
+        for row in rows.iterrows():
+            row = row[1]
+            class_ = self.classes_by_task[LabelType.BOUNDINGBOX].index(row["class"])
+            det = json.loads(row["value"])
+            box = np.array([class_, det[0], det[1], det[2], det[3]]).reshape(1, 5)
+            boxes = np.append(boxes, box, axis=0)
+        return boxes
+
+    def _load_polyline(self, rows: pd.DataFrame, height: int, width: int) -> np.ndarray:
+        seg = np.zeros((self.ns, height, width))
+        for row in rows.iterrows():
+            row = row[1]
+            class_ = self.classes_by_task[LabelType.SEGMENTATION].index(row["class"])
+            polyline = json.loads(row["value"])
+            polyline = [
+                (round(coord[0] * width), round(coord[1] * height))
+                for coord in polyline
+            ]
+            mask = Image.new("L", (width, height), 0)
+            draw = ImageDraw.Draw(mask)
+            draw.polygon(polyline, fill=1, outline=1)
+            mask = np.array(mask)
+            seg[class_, ...] = seg[class_, ...] + mask
+        seg[seg > 0] = 1
+        return seg
+
+    def _load_segmentation(
+        self, rows: pd.DataFrame, height: int, width: int
+    ) -> np.ndarray:
+        seg = np.zeros((self.ns, height, width))
+        for row in rows.iterrows():
+            row = row[1]
+            class_ = self.classes_by_task[LabelType.SEGMENTATION].index(row["class"])
+            height, width, counts_str = json.loads(row["value"])
+            mask = mask_util.decode(
+                {"counts": counts_str.encode("utf-8"), "size": [height, width]}
             )
-
-        return img, annotations
-
-    def _load_image_with_annotations(self, idx: int) -> Tuple[np.ndarray, Dict]:
+            seg[class_, ...] = seg[class_, ...] + mask
+        seg[seg > 0] = 1
+        return seg
+
+    def _load_keypoints(self, rows: pd.DataFrame) -> np.ndarray:
+        keypoints = np.zeros((0, self.max_nk * 3 + 1))
+        for row in rows.iterrows():
+            row = row[1]
+            class_ = self.classes_by_task[LabelType.KEYPOINT].index(row["class"])
+            kps = np.array(json.loads(row["value"])).reshape((-1, 3)).astype(np.float32)
+            kps = kps.flatten()
+            nk = len(kps)
+            kps = np.concatenate([[class_], kps])
+            points = np.zeros((1, self.max_nk * 3 + 1))
+            points[0, : nk + 1] = kps
+            keypoints = np.append(keypoints, points, axis=0)
+        return keypoints
+
+    def _load_image_with_annotations(
+        self, idx: int
+    ) -> Tuple[np.ndarray, Dict[str, dict]]:
         """Loads image and its annotations based on index.
 
         @type idx: int
@@ -203,88 +321,44 @@ def _load_image_with_annotations(self, idx: int) -> Tuple[np.ndarray, Dict]:
 
         img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
 
-        ih, iw, _ = img.shape
+        height, width, _ = img.shape
         annotations = {}
 
         if sub_df.ndim == 1:
             sub_df = pd.DataFrame([sub_df])
 
-        classification_rows = sub_df[sub_df["type"] == "classification"]
-        box_rows = sub_df[sub_df["type"] == "box"]
-        segmentation_rows = sub_df[sub_df["type"] == "segmentation"]
-        polyline_rows = sub_df[sub_df["type"] == "polyline"]
-        keypoints_rows = sub_df[sub_df["type"] == "keypoints"]
+        task_groups = sub_df["task_group"].unique()
 
-        seg = np.zeros((self.ns, ih, iw))
+        for task_group in task_groups:
+            group_df = sub_df[sub_df["task_group"] == task_group]
+            _annotations = {}
+            classification_rows = group_df[sub_df["type"] == "classification"]
+            box_rows = group_df[sub_df["type"] == "box"]
+            segmentation_rows = group_df[sub_df["type"] == "segmentation"]
+            polyline_rows = group_df[sub_df["type"] == "polyline"]
+            keypoints_rows = group_df[sub_df["type"] == "keypoints"]
 
-        if len(classification_rows):
-            classes = [
-                row[1]["class"]
-                for row in classification_rows.iterrows()
-                if bool(row[1]["value"])
-            ]
-            classify = np.zeros(self.nc)
-            for cls in classes:
-                cls = self.classes_by_task[LabelType.CLASSIFICATION].index(cls)
-                classify[cls] = classify[cls] + 1
-            classify[classify > 0] = 1
-            annotations[LabelType.CLASSIFICATION] = classify
-
-        if len(box_rows):
-            boxes = np.zeros((0, 5))
-            for row in box_rows.iterrows():
-                row = row[1]
-                cls = self.classes_by_task[LabelType.BOUNDINGBOX].index(row["class"])
-                det = json.loads(row["value"])
-                box = np.array([cls, det[0], det[1], det[2], det[3]]).reshape(1, 5)
-                boxes = np.append(boxes, box, axis=0)
-            annotations[LabelType.BOUNDINGBOX] = boxes
-
-        if len(polyline_rows):
-            for row in polyline_rows.iterrows():
-                row = row[1]
-                cls = self.classes_by_task[LabelType.SEGMENTATION].index(row["class"])
-                polyline = json.loads(row["value"])
-                polyline = [
-                    (round(coord[0] * iw), round(coord[1] * ih)) for coord in polyline
-                ]
-                mask = Image.new("L", (iw, ih), 0)
-                draw = ImageDraw.Draw(mask)
-                draw.polygon(polyline, fill=1, outline=1)
-                mask = np.array(mask)
-                seg[cls, ...] = seg[cls, ...] + mask
-            seg[seg > 0] = 1
-            annotations[LabelType.SEGMENTATION] = seg
-
-        if len(segmentation_rows):
-            for row in segmentation_rows.iterrows():
-                row = row[1]
-                cls = self.classes_by_task[LabelType.SEGMENTATION].index(row["class"])
-                height, width, counts_str = json.loads(row["value"])
-                mask = mask_util.decode(
-                    {"counts": counts_str.encode("utf-8"), "size": [height, width]}
+            if not classification_rows.empty:
+                _annotations[LabelType.CLASSIFICATION] = self._load_classification(
+                    classification_rows
                 )
-                seg[cls, ...] = seg[cls, ...] + mask
-            seg[seg > 0] = 1
-            annotations[LabelType.SEGMENTATION] = seg
-
-        if len(keypoints_rows):
-            # TODO: test with multi-class keypoint instances where nk's are not equal
-            keypoints = np.zeros((0, self.max_nk * 3 + 1))
-            for row in keypoints_rows.iterrows():
-                row = row[1]
-                cls = self.classes_by_task[LabelType.KEYPOINT].index(row["class"])
-                kps = (
-                    np.array(json.loads(row["value"]))
-                    .reshape((-1, 3))
-                    .astype(np.float32)
+
+            if not box_rows.empty:
+                _annotations[LabelType.BOUNDINGBOX] = self._load_box(box_rows)
+
+            if not polyline_rows.empty:
+                _annotations[LabelType.SEGMENTATION] = self._load_polyline(
+                    polyline_rows, height, width
+                )
+
+            if not segmentation_rows.empty:
+                _annotations[LabelType.SEGMENTATION] = self._load_segmentation(
+                    segmentation_rows, height, width
                 )
-                kps = kps.flatten()
-                nk = len(kps)
-                kps = np.concatenate([[cls], kps])
-                points = np.zeros((1, self.max_nk * 3 + 1))
-                points[0, : nk + 1] = kps
-                keypoints = np.append(keypoints, points, axis=0)
-            annotations[LabelType.KEYPOINT] = keypoints
+
+            if not keypoints_rows.empty:
+                _annotations[LabelType.KEYPOINT] = self._load_keypoints(keypoints_rows)
+
+            annotations[task_group] = _annotations
 
         return img, annotations
diff --git a/luxonis_ml/data/parsers/base_parser.py b/luxonis_ml/data/parsers/base_parser.py
index 905ff69a..bd9392f6 100644
--- a/luxonis_ml/data/parsers/base_parser.py
+++ b/luxonis_ml/data/parsers/base_parser.py
@@ -3,9 +3,9 @@
 from pathlib import Path
 from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple
 
-from luxonis_ml.data import DatasetGeneratorFunction, LuxonisDataset
+from luxonis_ml.data import DatasetIterator, LuxonisDataset
 
-ParserOutput = Tuple[DatasetGeneratorFunction, List[str], Dict[str, Dict], List[str]]
+ParserOutput = Tuple[DatasetIterator, List[str], Dict[str, Dict], List[str]]
 """Type alias for parser output.
 
 Contains a function to create the annotation generator, list of classes names, skeleton
@@ -141,15 +141,15 @@ def parse_dir(self, dataset_dir: Path, **kwargs) -> LuxonisDataset:
         return self.dataset
 
     @staticmethod
-    def _get_added_images(generator: DatasetGeneratorFunction) -> List[str]:
+    def _get_added_images(generator: DatasetIterator) -> List[str]:
         """Returns list of unique images added by the generator function.
 
-        @type generator: L{DatasetGeneratorFunction}
+        @type generator: L{DatasetGenerator}
         @param generator: Generator function
         @rtype: List[str]
         @return: List of added images by generator function
         """
-        return list(set(item["file"] for item in generator()))
+        return list(set(item["file"] for item in generator))
 
     @staticmethod
     def _compare_stem_files(list1: Iterable[Path], list2: Iterable[Path]) -> bool:
diff --git a/luxonis_ml/data/parsers/classification_directory_parser.py b/luxonis_ml/data/parsers/classification_directory_parser.py
index c46a9bf3..08c8e3b8 100644
--- a/luxonis_ml/data/parsers/classification_directory_parser.py
+++ b/luxonis_ml/data/parsers/classification_directory_parser.py
@@ -1,7 +1,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
-from luxonis_ml.data import DatasetGenerator
+from luxonis_ml.data import DatasetIterator
 
 from .base_parser import BaseParser, ParserOutput
 
@@ -68,7 +68,7 @@ def from_split(self, class_dir: Path) -> ParserOutput:
         """
         class_names = [d.name for d in class_dir.iterdir() if d.is_dir()]
 
-        def generator() -> DatasetGenerator:
+        def generator() -> DatasetIterator:
             for class_name in class_names:
                 for img_path in (class_dir / class_name).iterdir():
                     yield {
@@ -78,6 +78,6 @@ def generator() -> DatasetGenerator:
                         "value": True,
                     }
 
-        added_images = self._get_added_images(generator)
+        added_images = self._get_added_images(generator())
 
-        return generator, class_names, {}, added_images
+        return generator(), class_names, {}, added_images
diff --git a/luxonis_ml/data/parsers/coco_parser.py b/luxonis_ml/data/parsers/coco_parser.py
index 70d13cdb..9cc94d34 100644
--- a/luxonis_ml/data/parsers/coco_parser.py
+++ b/luxonis_ml/data/parsers/coco_parser.py
@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from luxonis_ml.data import DatasetGenerator
+from luxonis_ml.data import DatasetIterator
 
 from .base_parser import BaseParser, ParserOutput
 
@@ -139,7 +139,7 @@ def from_split(self, image_dir: Path, annotation_path: Path) -> ParserOutput:
                     "edges": (np.array(cat["skeleton"]) - 1).tolist(),
                 }
 
-        def generator() -> DatasetGenerator:
+        def generator() -> DatasetIterator:
             for img in coco_images:
                 img_id = img["id"]
 
@@ -203,6 +203,6 @@ def generator() -> DatasetGenerator:
                             "value": keypoints,
                         }
 
-        added_images = self._get_added_images(generator)
+        added_images = self._get_added_images(generator())
 
-        return generator, class_names, skeletons, added_images
+        return generator(), class_names, skeletons, added_images
diff --git a/luxonis_ml/data/parsers/create_ml_parser.py b/luxonis_ml/data/parsers/create_ml_parser.py
index d7d7db8f..18dbd15d 100644
--- a/luxonis_ml/data/parsers/create_ml_parser.py
+++ b/luxonis_ml/data/parsers/create_ml_parser.py
@@ -4,7 +4,7 @@
 
 from PIL import Image
 
-from luxonis_ml.data import DatasetGenerator
+from luxonis_ml.data import DatasetIterator
 
 from .base_parser import BaseParser, ParserOutput
 
@@ -105,7 +105,7 @@ def from_split(self, image_dir: Path, annotation_path: Path) -> ParserOutput:
                 curr_annotations["bboxes"].append((class_name, bbox_xywh))
             images_annotations.append(curr_annotations)
 
-        def generator() -> DatasetGenerator:
+        def generator() -> DatasetIterator:
             for curr_annotations in images_annotations:
                 path = curr_annotations["path"]
                 for class_name in curr_annotations["classes"]:
@@ -123,6 +123,6 @@ def generator() -> DatasetGenerator:
                         "value": tuple(bbox),
                     }
 
-        added_images = self._get_added_images(generator)
+        added_images = self._get_added_images(generator())
 
-        return generator, list(class_names), {}, added_images
+        return generator(), list(class_names), {}, added_images
diff --git a/luxonis_ml/data/parsers/darknet_parser.py b/luxonis_ml/data/parsers/darknet_parser.py
index 5eae24bd..4852f3c1 100644
--- a/luxonis_ml/data/parsers/darknet_parser.py
+++ b/luxonis_ml/data/parsers/darknet_parser.py
@@ -1,7 +1,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
-from luxonis_ml.data import DatasetGenerator
+from luxonis_ml.data import DatasetIterator
 
 from .base_parser import BaseParser, ParserOutput
 
@@ -74,7 +74,7 @@ def from_split(self, image_dir: Path, classes_path: Path) -> ParserOutput:
         with open(classes_path) as f:
             class_names = {i: line.rstrip() for i, line in enumerate(f.readlines())}
 
-        def generator() -> DatasetGenerator:
+        def generator() -> DatasetIterator:
             for img_path in self._list_images(image_dir):
                 ann_path = img_path.with_suffix(".txt")
                 file = str(img_path)
@@ -108,6 +108,6 @@ def generator() -> DatasetGenerator:
                         "value": bbox_xywh,
                     }
 
-        added_images = self._get_added_images(generator)
+        added_images = self._get_added_images(generator())
 
-        return generator, list(class_names.values()), {}, added_images
+        return generator(), list(class_names.values()), {}, added_images
diff --git a/luxonis_ml/data/parsers/segmentation_mask_directory_parser.py b/luxonis_ml/data/parsers/segmentation_mask_directory_parser.py
index 99275270..0a673000 100644
--- a/luxonis_ml/data/parsers/segmentation_mask_directory_parser.py
+++ b/luxonis_ml/data/parsers/segmentation_mask_directory_parser.py
@@ -6,7 +6,7 @@
 import pandas as pd
 import pycocotools.mask as mask_util
 
-from luxonis_ml.data import DatasetGenerator
+from luxonis_ml.data import DatasetIterator
 
 from .base_parser import BaseParser, ParserOutput
 
@@ -98,7 +98,7 @@ def from_split(
 
         class_names = pd.Series(df[idx_class].values, index=df[idx_pixel_val]).to_dict()
 
-        def generator() -> DatasetGenerator:
+        def generator() -> DatasetIterator:
             for mask_path in seg_dir.glob("*_mask.*"):
                 image_path = next(image_dir.glob(f"{mask_path.stem[:-5]}.*"))
                 file = str(image_path.absolute())
@@ -132,5 +132,5 @@ def generator() -> DatasetGenerator:
                         "value": value,
                     }
 
-        added_images = self._get_added_images(generator)
-        return generator, list(class_names.values()), {}, added_images
+        added_images = self._get_added_images(generator())
+        return generator(), list(class_names.values()), {}, added_images
diff --git a/luxonis_ml/data/parsers/solo_parser.py b/luxonis_ml/data/parsers/solo_parser.py
index e61eaf51..ac8d83a5 100644
--- a/luxonis_ml/data/parsers/solo_parser.py
+++ b/luxonis_ml/data/parsers/solo_parser.py
@@ -8,7 +8,7 @@
 import numpy as np
 import pycocotools.mask as mask_util
 
-from luxonis_ml.data import DatasetGenerator
+from luxonis_ml.data import DatasetIterator
 
 from .base_parser import BaseParser, ParserOutput
 
@@ -147,7 +147,7 @@ def from_split(
         # if NOT, set them manually with LuxonisDataset.set_skeletons() as SOLO format does not
         # encode which keypoint_name belongs to which class
 
-        def generator() -> DatasetGenerator:
+        def generator() -> DatasetIterator:
             for dir_path, dir_names, _ in os.walk(split_path):
                 for dir_name in dir_names:
                     if not dir_name.startswith("sequence"):
@@ -254,10 +254,10 @@ def generator() -> DatasetGenerator:
                                         "value": keypoints,
                                     }
 
-        added_images = self._get_added_images(generator)
+        added_images = self._get_added_images(generator())
 
         return (
-            generator,
+            generator(),
             class_names,
             skeletons,
             added_images,
diff --git a/luxonis_ml/data/parsers/tensorflow_csv_parser.py b/luxonis_ml/data/parsers/tensorflow_csv_parser.py
index 81551755..e1b0b499 100644
--- a/luxonis_ml/data/parsers/tensorflow_csv_parser.py
+++ b/luxonis_ml/data/parsers/tensorflow_csv_parser.py
@@ -4,7 +4,7 @@
 import numpy as np
 import pandas as pd
 
-from luxonis_ml.data import DatasetGenerator
+from luxonis_ml.data import DatasetIterator
 
 from .base_parser import BaseParser, ParserOutput
 
@@ -101,7 +101,7 @@ def from_split(self, image_dir: Path, annotation_path: Path) -> ParserOutput:
             bbox_xywh = bbox_xywh.tolist()
             images_annotations[path]["bboxes"].append((class_name, bbox_xywh))
 
-        def generator() -> DatasetGenerator:
+        def generator() -> DatasetIterator:
             for path in images_annotations:
                 curr_annotations = images_annotations[path]
                 for class_name in curr_annotations["classes"]:
@@ -119,6 +119,6 @@ def generator() -> DatasetGenerator:
                         "value": tuple(bbox),
                     }
 
-        added_images = self._get_added_images(generator)
+        added_images = self._get_added_images(generator())
 
-        return generator, list(class_names), {}, added_images
+        return generator(), list(class_names), {}, added_images
diff --git a/luxonis_ml/data/parsers/voc_parser.py b/luxonis_ml/data/parsers/voc_parser.py
index 79feaf81..0b888d96 100644
--- a/luxonis_ml/data/parsers/voc_parser.py
+++ b/luxonis_ml/data/parsers/voc_parser.py
@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from luxonis_ml.data import DatasetGenerator
+from luxonis_ml.data import DatasetIterator
 
 from .base_parser import BaseParser, ParserOutput
 
@@ -116,7 +116,7 @@ def from_split(
                     curr_annotations["bboxes"].append((class_name, bbox_xywh))
             images_annotations.append(curr_annotations)
 
-        def generator() -> DatasetGenerator:
+        def generator() -> DatasetIterator:
             for curr_annotations in images_annotations:
                 path = str(curr_annotations["path"])
                 for class_name in curr_annotations["classes"]:
@@ -134,9 +134,9 @@ def generator() -> DatasetGenerator:
                         "value": tuple(bbox),
                     }
 
-        added_images = self._get_added_images(generator)
+        added_images = self._get_added_images(generator())
 
-        return generator, list(class_names), {}, added_images
+        return generator(), list(class_names), {}, added_images
 
     @staticmethod
     def _xml_find(root: ET.Element, tag: str) -> str:
diff --git a/luxonis_ml/data/parsers/yolov4_parser.py b/luxonis_ml/data/parsers/yolov4_parser.py
index bc04905e..07a5462f 100644
--- a/luxonis_ml/data/parsers/yolov4_parser.py
+++ b/luxonis_ml/data/parsers/yolov4_parser.py
@@ -3,7 +3,7 @@
 
 from PIL import Image
 
-from luxonis_ml.data import DatasetGenerator
+from luxonis_ml.data import DatasetIterator
 
 from .base_parser import BaseParser, ParserOutput
 
@@ -86,7 +86,7 @@ def from_split(
         with open(classes_path) as f:
             class_names = {i: line.rstrip() for i, line in enumerate(f.readlines())}
 
-        def generator() -> DatasetGenerator:
+        def generator() -> DatasetIterator:
             with open(annotation_path) as f:
                 annotation_data = [line.rstrip() for line in f.readlines()]
 
@@ -127,6 +127,6 @@ def generator() -> DatasetGenerator:
                         "value": tuple(bbox_xywh),
                     }
 
-        added_images = self._get_added_images(generator)
+        added_images = self._get_added_images(generator())
 
-        return generator, list(class_names.values()), {}, added_images
+        return generator(), list(class_names.values()), {}, added_images
diff --git a/luxonis_ml/data/parsers/yolov6_parser.py b/luxonis_ml/data/parsers/yolov6_parser.py
index 3e34373c..d0df993f 100644
--- a/luxonis_ml/data/parsers/yolov6_parser.py
+++ b/luxonis_ml/data/parsers/yolov6_parser.py
@@ -3,7 +3,7 @@
 
 import yaml
 
-from luxonis_ml.data import DatasetGenerator
+from luxonis_ml.data import DatasetIterator
 
 from .base_parser import BaseParser, ParserOutput
 
@@ -110,7 +110,7 @@ def from_split(
             i: class_name for i, class_name in enumerate(classes_data["names"])
         }
 
-        def generator() -> DatasetGenerator:
+        def generator() -> DatasetIterator:
             for img_path in self._list_images(image_dir):
                 ann_path = annotation_dir / img_path.with_suffix(".txt").name
 
@@ -142,6 +142,6 @@ def generator() -> DatasetGenerator:
                         "value": bbox_xywh,
                     }
 
-        added_images = self._get_added_images(generator)
+        added_images = self._get_added_images(generator())
 
-        return generator, list(class_names.values()), {}, added_images
+        return generator(), list(class_names.values()), {}, added_images
diff --git a/luxonis_ml/data/utils/data_utils.py b/luxonis_ml/data/utils/data_utils.py
index fab81652..e47dd1ae 100644
--- a/luxonis_ml/data/utils/data_utils.py
+++ b/luxonis_ml/data/utils/data_utils.py
@@ -1,86 +1,45 @@
 import json
-import typing
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any, List, Tuple, Union
 
-import cv2
 import numpy as np
 import pycocotools.mask as mask_util
-from typeguard import TypeCheckError, check_type
+from typing_extensions import TypeAlias
 
-from .constants import ANNOTATIONS_SCHEMA as schema
+# NOTE: This could be defined elsewhere
+Number: TypeAlias = Union[int, float]
+"""Alias for a number type, which can be either an integer or a float."""
 
+ClassificationType: TypeAlias = bool
+"""Alias for a classification type, which is a boolean."""
 
-def check_annotation(data: Dict) -> None:
-    """Checks whether annotations match the expected format. Throws an exception if
-    there is a formatting error.
+LabelType: TypeAlias = Union[str, int, float, bool]
+"""Alias for a custom label type, which can be a string, integer, float, or boolean."""
 
-    @type data: Dict
-    @param data: A dictionary representing annotations, mapping annotation types to
-        values.
-    @rtype: NoneType
-    @return: None
-    """
+BoxType: TypeAlias = Tuple[Number, Number, Number, Number]
+"""Alias for a bounding box type, which is a tuple of (x, y, width, height)."""
 
-    if len(schema.keys()) != len(data.keys()) or set(schema.keys()) != set(data.keys()):
-        raise Exception(
-            f"Given keys {data.keys()} do not match annotations schema {schema.keys()}"
-        )
-
-    for key in schema:
-        origin = typing.get_origin(schema[key])
-        if origin is not None:
-            if origin == Union:
-                check = False
-                for typ in typing.get_args(schema[key]):
-                    if isinstance(data[key], typ):
-                        check = True
-                if not check:
-                    raise Exception(
-                        f"Found type {type(data[key])} for key '{key}' but expected: {typing.get_args(schema[key])}"
-                    )
-            else:
-                raise NotImplementedError
-        elif not isinstance(data[key], schema[key]):
-            raise Exception(
-                f"Found type {type(data[key])} for key '{key}' but expected: {schema[key]}"
-            )
+SegmentationType: TypeAlias = Tuple[int, int, Union[List[int], bytes]]
+"""Alias for a segmentation type in the RLE format, which is a tuple of (height, width,
+counts).
+
+`counts` can be encoded bytes or a raw list of integers.
+"""
 
-    typ = data["type"]
-    value = data["value"]
-
-    if typ == "classification":
-        _check_value_type(typ, value, bool)
-    elif typ == "label":
-        _check_value_type(typ, value, Union[str, int, float, bool])
-    elif typ == "box":
-        _check_value_type(
-            typ,
-            value,
-            Tuple[
-                Union[int, float],
-                Union[int, float],
-                Union[int, float],
-                Union[int, float],
-            ],
-        )
-    elif typ == "polyline":
-        _check_value_type(
-            typ,
-            value,
-            List[Tuple[Union[int, float], Union[int, float]]],
-        )
-    elif typ == "segmentation":
-        _check_value_type(
-            typ,
-            value,
-            Tuple[int, int, Union[List[int], bytes]],
-        )
-    elif typ == "keypoints":
-        _check_value_type(
-            typ,
-            value,
-            List[Tuple[Union[int, float], Union[int, float], int]],
-        )
+PolylineType: TypeAlias = List[Tuple[Number, Number]]
+"""Alias for a polyline type, which is a list of (x, y) points."""
+
+KeypointsType: TypeAlias = List[Tuple[Number, Number, int]]
+"""Alias for a keypoints type, which is a list of (x, y, visibility) points.
+Visibility is an integer value of the following format:
+
+    - 0: Not visible
+    - 1: Occluded
+    - 2: Visible
+
+"""
+
+ArrayType: TypeAlias = str
+"""Alias for an array type, which is a path to a numpy array (.npy)."""
 
 
 def check_arrays(values: List[Any]) -> None:
@@ -93,49 +52,28 @@ def check_arrays(values: List[Any]) -> None:
     @return: None
     """
 
-    for value in values:
-        _check_array(value)
+    def _check_valid_array(path: str) -> bool:
+        try:
+            np.load(path)
+            return True
+        except Exception:
+            return False
 
-
-def _check_value_type(name: str, value: Any, typ: Any) -> None:
-    """Checks if a value is of a given type, and raises a TypeError if not."""
-    try:
-        check_type(value, typ)
-    except TypeCheckError as e:
-        raise TypeError(f"Value {value} for key {name} is not of type {typ}") from e
-
-
-def _check_array(value: Any) -> None:
-    if not isinstance(value, str):
-        raise Exception(f"Array value {value} must be a path to a numpy array (.npy)")
-    if not _check_valid_array(value):
-        raise Exception(f"Array at path {value} is not a valid numpy array (.npy)")
-
-
-def _check_valid_image(path: str) -> bool:
-    try:
-        image = cv2.imread(path)
-        return image is not None
-    except Exception:
-        return False
-
-
-def _check_valid_array(path: str) -> bool:
-    try:
-        np.load(path)
-        return True
-    except Exception:
-        return False
+    for value in values:
+        if not isinstance(value, str):
+            raise Exception(
+                f"Array value {value} must be a path to a numpy array (.npy)"
+            )
+        if not _check_valid_array(value):
+            raise Exception(f"Array at path {value} is not a valid numpy array (.npy)")
 
 
-def transform_segmentation_value(
-    value: Tuple[int, int, Union[bytes, List[int]]],
-) -> str:
+def transform_segmentation_value(value: SegmentationType) -> str:
     """Transforms a segmentation in RLE format to the format stored by LuxonisDataset.
     The format recognized by LuxonisDataset is still RLE, but a dumped JSON string of
     height, width, and compressed counts.
 
-    @type value: Tuple[int, int, Union[bytes, List[int]]]
+    @type value: L{SegmentationType}
     @param value: The segmentation value in RLE format of (height, width, counts).
         Counts can be encoded bytes or a raw list.
     @rtype: str
diff --git a/luxonis_ml/enums/__init__.py b/luxonis_ml/enums/__init__.py
index 513b5f41..c6cbd9d7 100644
--- a/luxonis_ml/enums/__init__.py
+++ b/luxonis_ml/enums/__init__.py
@@ -1,3 +1,3 @@
-from .enums import DatasetType, LabelType, SplitType
+from .enums import AnnotationType, DatasetType, LabelType, SplitType
 
-__all__ = ["LabelType", "DatasetType", "SplitType"]
+__all__ = ["LabelType", "DatasetType", "SplitType", "AnnotationType"]
diff --git a/luxonis_ml/enums/enums.py b/luxonis_ml/enums/enums.py
index 02e956d0..deb398e8 100644
--- a/luxonis_ml/enums/enums.py
+++ b/luxonis_ml/enums/enums.py
@@ -1,6 +1,16 @@
 from enum import Enum
 
 
+class AnnotationType(str, Enum):
+    CLASSIFICATION = "classification"
+    BOX = "box"
+    SEGMENTATION = "segmentation"
+    POLYLINE = "polyline"
+    KEYPOINTS = "keypoints"
+    LABEL = "label"
+    ARRAY = "array"
+
+
 class LabelType(str, Enum):
     CLASSIFICATION = "class"
     SEGMENTATION = "segmentation"
diff --git a/luxonis_ml/utils/filesystem.py b/luxonis_ml/utils/filesystem.py
index 5b20d010..30c1a504 100644
--- a/luxonis_ml/utils/filesystem.py
+++ b/luxonis_ml/utils/filesystem.py
@@ -468,7 +468,7 @@ def get_file_uuid(self, path: PathType, local: bool = False) -> str:
         return file_hash_uuid
 
     def get_file_uuids(
-        self, paths: List[PathType], local: bool = False
+        self, paths: Sequence[PathType], local: bool = False
     ) -> Dict[str, str]:
         """Computes the UUIDs for all files stored in the filesystem.
 
diff --git a/luxonis_ml/utils/requirements.txt b/luxonis_ml/utils/requirements.txt
index 2f219207..c7e8f56d 100644
--- a/luxonis_ml/utils/requirements.txt
+++ b/luxonis_ml/utils/requirements.txt
@@ -3,6 +3,7 @@ pydantic-settings>=2.1.0
 PyYAML>=6.0
 fsspec>=2023.1.0
 rich>=13.6.0
+typer>=0.12.0
 # The following packages are installed dynamically when required
 # s3fs>=2023.1.0
 # gcsfs>=2023.1.0


Test patch:
diff --git a/tests/test_data/test_parsers.py b/tests/test_data/test_parsers.py
index b6f903ee..e178466d 100644
--- a/tests/test_data/test_parsers.py
+++ b/tests/test_data/test_parsers.py
@@ -92,6 +92,6 @@ def test_dir_parser(
     assert len(dataset) > 0
     loader = LuxonisLoader(dataset)
     _, ann = next(iter(loader))
-    tasks = ann.keys()
+    tasks = ann["default"].keys()
     assert set(tasks) == set(expected_tasks)
     dataset.delete_dataset()