Move get_idx_list_per_classes to dataset class and address other PR comments.

Signed-off-by: Albert van Houten <albert.van.houten@intel.com>

Author: AlbertvanHouten

lib/src/otx/data/dataset/base.py

@@ -6,13 +6,14 @@
 from __future__ import annotations
 
 from abc import abstractmethod
+from collections import defaultdict
 from collections.abc import Iterable
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, Callable, Iterator, List, Union
 
 import cv2
 import numpy as np
-from datumaro.components.annotation import AnnotationType
+from datumaro.components.annotation import AnnotationType, LabelCategories
 from datumaro.util.image import IMAGE_BACKEND, IMAGE_COLOR_CHANNEL, ImageBackend
 from datumaro.util.image import ImageColorChannel as DatumaroImageColorChannel
 from torch.utils.data import Dataset
@@ -196,3 +197,18 @@ def _get_item_impl(self, idx: int) -> OTXDataItem | None:
     def collate_fn(self) -> Callable:
         """Collection function to collect KeypointDetDataEntity into KeypointDetBatchDataEntity in data loader."""
         return OTXDataItem.collate_fn
+
+    def get_idx_list_per_classes(self, use_string_label: bool = False) -> dict[int | str, list[int]]:
+        """Compute class statistics."""
+        stats: dict[int | str, list[int]] = defaultdict(list)
+        for item_idx, item in enumerate(self.dm_subset):
+            for ann in item.annotations:
+                if use_string_label:
+                    labels = self.dm_subset.categories().get(AnnotationType.label, LabelCategories())
+                    stats[labels.items[ann.label].name].append(item_idx)
+                else:
+                    stats[ann.label].append(item_idx)
+        # Remove duplicates in label stats idx: O(n)
+        for k in stats:
+            stats[k] = list(dict.fromkeys(stats[k]))
+        return stats

lib/src/otx/data/dataset/base_new.py

@@ -157,5 +157,5 @@ def collate_fn(self) -> Callable:
         return _default_collate_fn
 
     @abc.abstractmethod
-    def get_idx_list_per_classes(self) -> dict[int, list[int]]:
+    def get_idx_list_per_classes(self, use_string_label: bool = False) -> dict[int, list[int]]:
         """Get a dictionary with class labels as keys and lists of corresponding sample indices as values."""

lib/src/otx/data/dataset/classification_new.py

@@ -21,12 +21,14 @@ def __init__(self, **kwargs) -> None:
         kwargs["sample_type"] = ClassificationSample
         super().__init__(**kwargs)
 
-    def get_idx_list_per_classes(self) -> dict[int, list[int]]:
+    def get_idx_list_per_classes(self, use_string_label: bool = False) -> dict[int, list[int]]:
         """Get index list per class."""
         idx_list_per_classes: dict[int, list[int]] = {}
         for idx in range(len(self)):
             item = self.dm_subset[idx]
             label_id = item.label.item()
+            if use_string_label:
+                label_id = self.label_info.labels[label_id]
             if label_id not in idx_list_per_classes:
                 idx_list_per_classes[label_id] = []
             idx_list_per_classes[label_id].append(idx)

lib/src/otx/data/factory.py

@@ -82,13 +82,13 @@ def create(
                 categories = cls._get_label_categories(dm_subset, data_format)
                 dataset = DatasetNew(ClassificationSample, categories={"label": categories})
                 for item in dm_subset:
-                    if len(item.media.data.shape) == 3:  # TODO: Account for grayscale images
+                    if len(item.media.data.shape) == 3:  # TODO(albert): Account for grayscale images
                         dataset.append(ClassificationSample.from_dm_item(item))
                 common_kwargs["dm_subset"] = dataset
             return OTXMulticlassClsDataset(**common_kwargs)
 
         if task == OTXTaskType.MULTI_LABEL_CLS:
-            from otx.data.dataset.classification import OTXMultilabelClsDataset
+            from .dataset.classification import OTXMultilabelClsDataset
 
             return OTXMultilabelClsDataset(**common_kwargs)
 

lib/src/otx/data/samplers/balanced_sampler.py

@@ -9,12 +9,8 @@
 from typing import TYPE_CHECKING
 
 import torch
-from datumaro import DatasetSubset
-from datumaro.experimental import Dataset as NewDataset
 from torch.utils.data import Sampler
 
-from otx.data.utils import get_idx_list_per_classes
-
 if TYPE_CHECKING:
     from otx.data.dataset.base import OTXDataset
     from otx.data.dataset.base_new import OTXDataset as OTXDatasetNew
@@ -64,11 +60,7 @@ def __init__(
         super().__init__(dataset)
 
         # img_indices: dict[label: list[idx]]
-        ann_stats: dict[int | str, list[int]]
-        if isinstance(dataset.dm_subset, DatasetSubset):
-            ann_stats = get_idx_list_per_classes(dataset.dm_subset)
-        elif isinstance(dataset.dm_subset, NewDataset):
-            ann_stats = dataset.get_idx_list_per_classes()  # type: ignore[attr-defined]
+        ann_stats = dataset.get_idx_list_per_classes()
 
         self.img_indices = {k: torch.tensor(v, dtype=torch.int64) for k, v in ann_stats.items() if len(v) > 0}
         self.num_cls = len(self.img_indices.keys())

lib/src/otx/data/samplers/class_incremental_sampler.py

@@ -12,7 +12,6 @@
 from torch.utils.data import Sampler
 
 from otx.data.dataset.base import OTXDataset
-from otx.data.utils import get_idx_list_per_classes
 
 
 class ClassIncrementalSampler(Sampler):
@@ -65,7 +64,7 @@ def __init__(
         super().__init__(dataset)
 
         # Need to split new classes dataset indices & old classses dataset indices
-        ann_stats = get_idx_list_per_classes(dataset.dm_subset, True)
+        ann_stats = dataset.get_idx_list_per_classes(True)
         new_indices, old_indices = [], []
         for cls in new_classes:
             new_indices.extend(ann_stats[cls])

lib/src/otx/data/utils/__init__.py

@@ -7,7 +7,6 @@
     adapt_input_size_to_dataset,
     adapt_tile_config,
     get_adaptive_num_workers,
-    get_idx_list_per_classes,
     import_object_from_module,
     instantiate_sampler,
 )
@@ -17,6 +16,5 @@
     "adapt_input_size_to_dataset",
     "instantiate_sampler",
     "get_adaptive_num_workers",
-    "get_idx_list_per_classes",
     "import_object_from_module",
 ]

lib/src/otx/data/utils/utils.py

@@ -15,14 +15,13 @@
 import cv2
 import numpy as np
 import torch
-from datumaro.components.annotation import AnnotationType, Bbox, ExtractedMask, LabelCategories, Polygon
+from datumaro.components.annotation import AnnotationType, Bbox, ExtractedMask, Polygon
 from datumaro.components.annotation import Shape as _Shape
 
 from otx.types import OTXTaskType
 from otx.utils.device import is_xpu_available
 
 if TYPE_CHECKING:
-    from datumaro import Dataset as DmDataset
     from datumaro import DatasetSubset
     from torch.utils.data import Dataset, Sampler
 
@@ -322,22 +321,6 @@ def get_adaptive_num_workers(num_dataloader: int = 1) -> int | None:
     return min(cpu_count() // (num_dataloader * num_devices), 8)  # max available num_workers is 8
 
 
-def get_idx_list_per_classes(dm_dataset: DmDataset, use_string_label: bool = False) -> dict[int | str, list[int]]:
-    """Compute class statistics."""
-    stats: dict[int | str, list[int]] = defaultdict(list)
-    for item_idx, item in enumerate(dm_dataset):
-        for ann in item.annotations:
-            if use_string_label:
-                labels = dm_dataset.categories().get(AnnotationType.label, LabelCategories())
-                stats[labels.items[ann.label].name].append(item_idx)
-            else:
-                stats[ann.label].append(item_idx)
-    # Remove duplicates in label stats idx: O(n)
-    for k in stats:
-        stats[k] = list(dict.fromkeys(stats[k]))
-    return stats
-
-
 def import_object_from_module(obj_path: str) -> Any:  # noqa: ANN401
     """Get object from import format string."""
     module_name, obj_name = obj_path.rsplit(".", 1)

lib/tests/unit/data/conftest.py

@@ -41,7 +41,7 @@
 
 
 @pytest.fixture(params=["bytes", "file"])
-def fxt_dm_item(requeset, tmpdir) -> DatasetItem:
+def fxt_dm_item(request, tmpdir) -> DatasetItem:
     np_img = np.zeros(shape=(10, 10, 3), dtype=np.uint8)
     np_img[:, :, 0] = 0  # Set 0 for B channel
     np_img[:, :, 1] = 1  # Set 1 for G channel

lib/tests/unit/data/samplers/test_balanced_sampler.py

@@ -12,7 +12,6 @@
 
 from otx.data.dataset.base import OTXDataset
 from otx.data.samplers.balanced_sampler import BalancedSampler
-from otx.data.utils import get_idx_list_per_classes
 
 
 @pytest.fixture()
@@ -81,7 +80,7 @@ def test_sampler_iter_with_multiple_replicas(self, fxt_imbalanced_dataset):
 
     def test_compute_class_statistics(self, fxt_imbalanced_dataset):
         # Compute class statistics
-        stats = get_idx_list_per_classes(fxt_imbalanced_dataset.dm_subset)
+        stats = fxt_imbalanced_dataset.get_idx_list_per_classes()
 
         # Check the expected results
         assert stats == {0: list(range(100)), 1: list(range(100, 108))}
@@ -90,7 +89,7 @@ def test_sampler_iter_per_class(self, fxt_imbalanced_dataset):
         batch_size = 4
         sampler = BalancedSampler(fxt_imbalanced_dataset)
 
-        stats = get_idx_list_per_classes(fxt_imbalanced_dataset.dm_subset)
+        stats = fxt_imbalanced_dataset.get_idx_list_per_classes()
         class_0_idx = stats[0]
         class_1_idx = stats[1]
         list_iter = list(iter(sampler))