Custom Loaders Support

[kozlov721]

Added support for using custom loaders.

Example

source.py

from luxonis_train.utils.loaders import BaseLoaderTorch

class CustomLoader(BaseLoaderTorch):
    import torch
    from luxonis_ml.data import LabelType

    from luxonis_train.utils.loaders import LuxonisLoaderTorchOutput

    @property
    def input_shape(self) -> torch.Size:
        return self.torch.Size((1, 3, 256, 256))

    def __getitem__(self, _: int) -> LuxonisLoaderTorchOutput:
        from luxonis_ml.data import LabelType

        img = self.torch.rand(3, 256, 256)
        label = self.torch.zeros(1, 256, 256)
        label[0, 100:200, 100:200] = 1
        labels = {
            "default": {
                LabelType.SEGMENTATION: label,
            }
        }
        return img, labels

    def __len__(self) -> int:
        return 32

    def get_classes(self) -> dict[LabelType, list[str]]:
        return {self.LabelType.SEGMENTATION: ["square"]}

luxonis_train --source source.py \
              --config configs/segmentation_model.yaml \ 
              loader.name CustomLoader loader.params '{}'

Breaking Config Changes

dataset section renamed to loader. Custom loader can be referenced by its name in the LOADERS registry. The default loader is LuxonisLoaderTorch. Former keys of the dataset field now go under params field in the loader.

loader:
  name: LuxonisLoaderTorch  # this is the default

  train_view: train
  val_view: val
  test_view: test

  params:
    dataset_name: coco_test
    bucket_storage: gcs

[github-actions]

Test Results

  6 files    6 suites   38m 16s ⏱️
 57 tests  25 ✅ 0 💤  7 ❌  25 🔥
342 runs  178 ✅ 0 💤 14 ❌ 150 🔥

For more details on these failures and errors, see this check.

Results for commit 49aab2b.

[github-actions]

☂️ Python Coverage

current status: ✅

Overall Coverage

Lines	Covered	Coverage	Threshold	Status
4713	3650	77%	0%	🟢

New Files

No new covered files...

Modified Files

File	Coverage	Status
luxonis_train/init.py	100%	🟢
luxonis_train/main.py	49%	🟢
luxonis_train/callbacks/test_on_train_end.py	86%	🟢
luxonis_train/core/init.py	100%	🟢
luxonis_train/core/archiver.py	19%	🟢
luxonis_train/core/core.py	78%	🟢
luxonis_train/core/exporter.py	55%	🟢
luxonis_train/core/inferer.py	87%	🟢
luxonis_train/core/trainer.py	62%	🟢
luxonis_train/core/tuner.py	83%	🟢
luxonis_train/models/luxonis_model.py	89%	🟢
luxonis_train/utils/config.py	95%	🟢
luxonis_train/utils/general.py	78%	🟢
luxonis_train/utils/loaders/init.py	100%	🟢
luxonis_train/utils/loaders/base_loader.py	63%	🟢
luxonis_train/utils/loaders/luxonis_loader_torch.py	100%	🟢
luxonis_train/utils/registry.py	100%	🟢
TOTAL	79%	🟢

updated for commit: 27e59fc by action🐍

[conorsim]

Looks like a nice improvement, thanks!

[conorsim]

Why do we skip this test?

[kozlov721]

The NN Archive in the latest dev of luxonis-ml is not compatible with the tests for now (we need to figure out the yolov8 keypoint outputs in the archive)

[tersekmatija]

1/ Does this affect any of our current online datasets?
2/ Have we tested the kpts training?

Otherwise looks good, thanks!

[tersekmatija]

What if self.loaders[view] does not exist?

[kozlov721]

It will always exist

[kozlov721]

1/ Does this affect any of our current online datasets?

No, this shouldn't change anything regarding the underlying datasets

2/ Have we tested the kpts training?

Yes

[CaptainTrojan]

@kozlov721 What is the correct way of providing BucketType and BucketStorage to the default LuxonisLoaderTorch? The DatasetConfig ensured that strings like internal, local, gcs etc. were mapped to the corresponding enum type.