Multi Input Support

[CaptainTrojan]

Motivation

The goal is to support multi-input models and loaders in luxonis-train.

Examples

Multi-input loader

The loader simulates an input consisting of stereo images (left and right), disparity maps, and point clouds. It also identifies "left" as the primary image type, which can be used in visualizers.

class CustomMultiInputLoader(BaseLoaderTorch):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @property
    def input_shape(self) -> dict[str, Size]:
        return {
            "left": torch.Size([3, 224, 224]),
            "right": torch.Size([3, 224, 224]),
            "disparity": torch.Size([1, 224, 224]),
            "pointcloud": torch.Size([1000, 3]),
        }

    @property
    def image_source(self) -> str:
        return "left"

    def __getitem__(self, idx):
        # Fake data
        left = torch.rand(3, 224, 224, dtype=torch.float32)
        right = torch.rand(3, 224, 224, dtype=torch.float32)
        disparity = torch.rand(1, 224, 224, dtype=torch.float32)
        pointcloud = torch.rand(1000, 3, dtype=torch.float32)
        inputs = {
            "left": left,
            "right": right,
            "disparity": disparity,
            "pointcloud": pointcloud,
        }

        # Fake labels
        segmap = torch.zeros(1, 224, 224, dtype=torch.float32)
        labels = {
            "segmentation": (segmap, LabelType.SEGMENTATION)
        }

        return inputs, labels

    def __len__(self):
        return 10

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

Multi-input model

This model includes nodes such as FullBackbone, RGBDBackbone (processing left, right, and disparity inputs), and PointcloudBackbone (processing pointcloud inputs). It also features FusionNeck and FusionNeck2 for combining outputs from these backbones, and CustomSegHead1 and CustomSegHead2 which also make use of the different fusion necks and inputs. The model showcases our capability to handle intricate data fusion, multi-modal processing, and serving loader data to any node in the graph.

model:
  name: example_multi_input
  nodes:
    - name: FullBackbone
      alias: full_backbone  # By default loads all loader sub-elements since no inputs nor input_sources are provided

    - name: RGBDBackbone
      alias: rgbd_backbone
      input_sources:
        - left
        - right
        - disparity

    - name: PointcloudBackbone
      alias: pointcloud_backbone
      input_sources:
        - pointcloud

    - name: FusionNeck
      alias: fusion_neck
      inputs:
        - rgbd_backbone
        - pointcloud_backbone
      input_sources:
        - disparity

    - name: FusionNeck2
      alias: fusion_neck_2
      inputs:
        - rgbd_backbone
        - pointcloud_backbone
        - full_backbone
    
    - name: CustomSegHead1
      alias: head_1
      inputs:
        - fusion_neck

    - name: CustomSegHead2
      alias: head_2
      inputs:
        - fusion_neck
        - fusion_neck_2
      input_sources:
        - disparity

Test changes

Updated collate_fn test due to how inputs were changed from Tensors to dict[str, Tensor].

Breaking changes

Changed the type of original_in_shape of LuxonisModel from Size to dict[str, Size] (multiple inputs), so any custom models making use of this will stop working. To resolve this, users can make use of the new structure like this.

[github-actions]

☂️ Python Coverage

current status: ✅

Overall Coverage

Lines	Covered	Coverage	Threshold	Status
4885	3758	77%	0%	🟢

New Files

No new covered files...

Modified Files

File	Coverage	Status
luxonis_train/attached_modules/base_attached_module.py	84%	🟢
luxonis_train/attached_modules/losses/adaptive_detection_loss.py	98%	🟢
luxonis_train/attached_modules/metrics/mean_average_precision.py	100%	🟢
luxonis_train/attached_modules/metrics/mean_average_precision_keypoints.py	94%	🟢
luxonis_train/attached_modules/metrics/object_keypoint_similarity.py	95%	🟢
luxonis_train/attached_modules/visualizers/utils.py	80%	🟢
luxonis_train/core/core.py	80%	🟢
luxonis_train/models/luxonis_model.py	90%	🟢
luxonis_train/models/luxonis_output.py	72%	🟢
luxonis_train/nodes/base_node.py	76%	🟢
luxonis_train/nodes/bisenet_head.py	50%	🟢
luxonis_train/nodes/classification_head.py	100%	🟢
luxonis_train/nodes/efficient_bbox_head.py	100%	🟢
luxonis_train/nodes/implicit_keypoint_bbox_head.py	92%	🟢
luxonis_train/nodes/segmentation_head.py	100%	🟢
luxonis_train/utils/boxutils.py	83%	🟢
luxonis_train/utils/config.py	95%	🟢
luxonis_train/utils/general.py	79%	🟢
luxonis_train/utils/loaders/base_loader.py	76%	🟢
luxonis_train/utils/loaders/luxonis_loader_torch.py	100%	🟢
TOTAL	87%	🟢

updated for commit: 9987d0d by action🐍

[CaptainTrojan]

I see that the MacOS tests fail for all open PRs, so opening this one as well.

[kozlov721]

I would rename this to maybe source_name. It doesn't have to necessarily be images

[kozlov721]

Also raising an error in case it's None.

[kozlov721]

I would rename this to maybe source_name. It doesn't have to necessarily be images

Or do we always expect this to be an image?

[CaptainTrojan]

The property images_name I created due to compatibility with image loggers/visualizers, which expect some of the loader inputs to contain image-like objects which it can process in OpenCV, log as images, etc. Right now, the config.yaml owner is responsible for setting this property for their custom multi-input logger so that the rest of the code knows which inputs are images.

[CaptainTrojan]

Also raising an error in case it's None.

I moved the declaration of images_name to the config, which has 'features' as the value by default, so unless the user implements their custom loader overriding this property as None, it should never be None. Is that how you produced this scenario?

[kozlov721]

I see, that makes sense. Though I'd maybe rename it to image_source or something like that.

[kozlov721]

I see that the MacOS tests fail for all open PRs, so opening this one as well.

Yeah, you can ignore the MacOS tests. The MacOS machines are running out of memory, I'm planning to fix it eventually, but it's low priority.

[github-actions]

Test Results

  4 files    4 suites   1h 0m 47s ⏱️
 68 tests  43 ✅  25 💤 0 ❌
272 runs  172 ✅ 100 💤 0 ❌

Results for commit 9987d0d.

♻️ This comment has been updated with latest results.