Full Custom Tasks Support

[kozlov721]

Added full support for custom task names. When using custom task names, a new field task has to be specified in the configuration of the corresponding head node. The task gets propagated to all the modules attached to the head (losses, metrics, etc.).

Example

Labels

labels = {"segmentation-task": ...,
          "boundingbox-task": ...,
          "keypoints-task": ...}

config.yaml

nodes:
    ...
    - name: SegmentationHead
      # single-task head requires just the new task name
      task: segmentation-task

    - name: EfficientKeypointBBoxHead
      # multi-task heads require a mapping from the default names 
      # (string representations of the corresponding `LabelType`) to the new custom task names
      task:
        keypoints: keypoints-task
        boundingbox: boundingbox-task

Subclassing Changes

Minor breaking changes introduced for implementing custom nodes and attached modules.

• Specifying task types and required labels using class attributes instead of __init__ arguments.
• Removed protocols for attached nodes
  • BaseAttachedNode.get_label and BaseAttachedNode.get_input_tensor now do the checking insteadd

Nodes

old_node.py

class CustomNode(BaseNode):
    def __init__(self, **kwargs):
        super().__init__(**kwargs, _task_type=LabelType.SEGMENTATION)

new_node.py

class CustomNode(BaseNode):
    # either a list of `LabelType`s or a mapping from `LabelType` to a default task name,
    # the following is equivalent to `tasks = {LabelType.SEGMENTATION: "segmentation"}`
    tasks = [LabelType.SEGMENTATION]
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

Attached Modules

old_loss.py

class CustomLoss(BaseLoss):
    def __init__(self, **kwargs):
        super().__init__(
            **kwargs, 
            required_labels=LabelType.SEGMENTATION,
            protocol=SegmentationProtocol,
        )
    def prepare(self, inputs: Packet[Tensor], labels: Labels):
        # presence of the "segmentation" key is validated beforehand by the `Protocol`
        predicted_masks = inputs["segmentation"]
        # validated beforehand using `required_labels`
        label_mask = inputs["segmentation"]

new_loss.py

class CustomLoss(BaseLoss):
    # list of all tasks this module supports
    # can contain either standalone `LabelType`s or tuples of `LabelType`s
    # tuple specifies that all of the tasks in the apple must be provided  
    supported_labels = [LabelType.SEGMENTATION]
    # [(LabelType.BOUNDINGBOX, LabelType.KEYPOINTS)] would mean that the module
    # requires both of these tasks to be present in node outputs and labels
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def prepare(self, inputs: Packet[Tensor], labels: Labels):
        # returns the correct input tensor from the node outputs
        predicted_masks = self.get_input_tensor(inputs, LabelType.SEGMENTATION)
        # returns the correct label based on the task of the attached `BaseNode`
        label_mask = self.get_label(labels, LabelType.SEGMENTATION)

[github-actions]

Test Results

  4 files    4 suites   1h 15m 40s ⏱️
 69 tests  45 ✅ 24 💤 0 ❌
276 runs  180 ✅ 96 💤 0 ❌

Results for commit 67752ed.

♻️ This comment has been updated with latest results.

[klemen1999]

LGTM!

Just two question:

• Now that most of the protocols were removed do we still keep the FeaturesProtocol?
• Are the LuxonisDataset datasets created before still compatible with this newer version of luxonis-train/luxonis-ml?

[kozlov721]

Now that most of the protocols were removed do we still keep the FeaturesProtocol?

That one is used in BaseNode subclasses which still use protocols. I'll probably remove it too at some point, but that would be in a separate PR.

Are the LuxonisDataset datasets created before still compatible with this newer version of luxonis-train/luxonis-ml?

No, only datasets created after #110 are compatible (luxonis-ml>=0.2.0).

[github-actions]

☂️ Python Coverage

current status: ✅

Overall Coverage

Lines	Covered	Coverage	Threshold	Status
5291	4152	78%	0%	🟢

New Files

No new covered files...

Modified Files

File	Coverage	Status
luxonis_train/main.py	73%	🟢
luxonis_train/attached_modules/base_attached_module.py	83%	🟢
luxonis_train/attached_modules/losses/adaptive_detection_loss.py	98%	🟢
luxonis_train/attached_modules/losses/base_loss.py	92%	🟢
luxonis_train/attached_modules/losses/bce_with_logits.py	93%	🟢
luxonis_train/attached_modules/losses/cross_entropy.py	96%	🟢
luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py	98%	🟢
luxonis_train/attached_modules/losses/implicit_keypoint_bbox_loss.py	97%	🟢
luxonis_train/attached_modules/losses/keypoint_loss.py	97%	🟢
luxonis_train/attached_modules/losses/sigmoid_focal_loss.py	60%	🟢
luxonis_train/attached_modules/losses/smooth_bce_with_logits.py	78%	🟢
luxonis_train/attached_modules/losses/softmax_focal_loss.py	36%	🟢
luxonis_train/attached_modules/metrics/base_metric.py	88%	🟢
luxonis_train/attached_modules/metrics/common.py	94%	🟢
luxonis_train/attached_modules/metrics/mean_average_precision.py	100%	🟢
luxonis_train/attached_modules/metrics/mean_average_precision_keypoints.py	95%	🟢
luxonis_train/attached_modules/metrics/object_keypoint_similarity.py	87%	🟢
luxonis_train/attached_modules/visualizers/base_visualizer.py	92%	🟢
luxonis_train/attached_modules/visualizers/bbox_visualizer.py	96%	🟢
luxonis_train/attached_modules/visualizers/classification_visualizer.py	94%	🟢
luxonis_train/attached_modules/visualizers/keypoint_visualizer.py	100%	🟢
luxonis_train/attached_modules/visualizers/multi_visualizer.py	88%	🟢
luxonis_train/attached_modules/visualizers/segmentation_visualizer.py	81%	🟢
luxonis_train/callbacks/luxonis_progress_bar.py	89%	🟢
luxonis_train/core/inferer.py	90%	🟢
luxonis_train/models/luxonis_model.py	90%	🟢
luxonis_train/models/predefined_models/classification_model.py	100%	🟢
luxonis_train/models/predefined_models/detection_model.py	100%	🟢
luxonis_train/models/predefined_models/keypoint_detection_model.py	95%	🟢
luxonis_train/models/predefined_models/segmentation_model.py	100%	🟢
luxonis_train/nodes/base_node.py	78%	🟢
luxonis_train/nodes/bisenet_head.py	52%	🟢
luxonis_train/nodes/classification_head.py	100%	🟢
luxonis_train/nodes/efficient_bbox_head.py	100%	🟢
luxonis_train/nodes/efficient_keypoint_bbox_head.py	99%	🟢
luxonis_train/nodes/implicit_keypoint_bbox_head.py	92%	🟢
luxonis_train/nodes/rexnetv1.py	15%	🟢
luxonis_train/nodes/segmentation_head.py	100%	🟢
luxonis_train/utils/assigners/atts_assigner.py	100%	🟢
luxonis_train/utils/config.py	95%	🟢
luxonis_train/utils/loaders/init.py	100%	🟢
luxonis_train/utils/loaders/luxonis_loader_torch.py	100%	🟢
luxonis_train/utils/types.py	83%	🟢
TOTAL	88%	🟢

updated for commit: 67752ed by action🐍

[conorsim]

I didn't review the code, but from manual testing some custom task applications it looks good