Sort classes alphabetically (#245)

Co-authored-by: Martin Kozlovsky <[email protected]>

Author: JSabadin

luxonis_ml/data/__init__.py

@@ -12,6 +12,7 @@
         LuxonisComponent,
         LuxonisDataset,
         LuxonisSource,
+        Metadata,
         UpdateMode,
     )
     from .loaders import LOADERS_REGISTRY, BaseLoader, LuxonisLoader
@@ -45,6 +46,7 @@ def load_loader_plugins() -> None:  # pragma: no cover
     "DatasetIterator",
     "DATASETS_REGISTRY",
     "Category",
+    "Metadata",
     "LOADERS_REGISTRY",
     "ImageType",
     "LuxonisComponent",

luxonis_ml/data/datasets/__init__.py

@@ -10,6 +10,7 @@
 )
 from .base_dataset import DATASETS_REGISTRY, BaseDataset, DatasetIterator
 from .luxonis_dataset import LuxonisDataset, UpdateMode
+from .metadata import Metadata
 from .source import LuxonisComponent, LuxonisSource
 
 __all__ = [
@@ -28,4 +29,5 @@
     "Detection",
     "ArrayAnnotation",
     "UpdateMode",
+    "Metadata",
 ]

luxonis_ml/data/datasets/base_dataset.py

@@ -1,5 +1,14 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Iterator, List, Optional, Sequence, Tuple, Type, Union
+from typing import (
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    Union,
+)
 
 from semver.version import Version
 from typing_extensions import TypeAlias
@@ -60,13 +69,19 @@ def update_source(self, source: LuxonisSource) -> None:
 
     @abstractmethod
     def set_classes(
-        self, classes: List[str], task: Optional[str] = None
+        self,
+        classes: Union[List[str], Dict[str, int]],
+        task: Optional[str] = None,
     ) -> None:
-        """Sets the names of classes for the dataset. This can be across
-        all CV tasks or certain tasks.
-
-        @type classes: List[str]
-        @param classes: List of class names to set.
+        """Sets the classes for the dataset. This can be across all CV
+        tasks or certain tasks.
+
+        @type classes: Union[List[str], Dict[str, int]]
+        @param classes: Either a list of class names or a dictionary
+            mapping class names to class IDs. If list is provided, the
+            class IDs will be assigned I{alphabetically} starting from
+            C{0}. If the class names contain the class C{"background"},
+            it will be assigned the class ID C{0}.
         @type task: Optional[str]
         @param task: Optionally specify the task where these classes
             apply.