Instance Segmentation Head

README.md

@@ -595,7 +595,7 @@ from luxonis_train import LuxonisLightningModule
 from luxonis_train.utils.registry import CALLBACKS
 
 
-@CALLBACKS.register_module()
+@CALLBACKS.register()
 class CustomCallback(pl.Callback):
     def __init__(self, message: str, **kwargs):
         super().__init__(**kwargs)

configs/README.md

@@ -280,14 +280,14 @@ We use [`Albumentations`](https://albumentations.ai/docs/) library for `augmenta
 
 Additionally, we support `Mosaic4` and `MixUp` batch augmentations and letterbox resizing if `keep_aspect_ratio: true`.
 
-| Key                 | Type         | Default value | Description                                                                                                                                                             |
-| ------------------- | ------------ | ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `train_image_size`  | `list[int]`  | `[256, 256]`  | Image size used for training as `[height, width]`                                                                                                                       |
-| `keep_aspect_ratio` | `bool`       | `True`        | Whether to keep the aspect ratio while resizing                                                                                                                         |
-| `train_rgb`         | `bool`       | `True`        | Whether to train on RGB or BGR images                                                                                                                                   |
-| `normalize.active`  | `bool`       | `True`        | Whether to use normalization                                                                                                                                            |
-| `normalize.params`  | `dict`       | `{}`          | Parameters for normalization, see [Normalize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize) |
-| `augmentations`     | `list[dict]` | `[]`          | List of `Albumentations` augmentations                                                                                                                                  |
+| Key                 | Type                    | Default value | Description                                                                                                                                                             |
+| ------------------- | ----------------------- | ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `train_image_size`  | `list[int]`             | `[256, 256]`  | Image size used for training as `[height, width]`                                                                                                                       |
+| `keep_aspect_ratio` | `bool`                  | `True`        | Whether to keep the aspect ratio while resizing                                                                                                                         |
+| `color_space`       | `Literal["RGB", "BGR"]` | `"RGB"`       | Whether to train on RGB or BGR images                                                                                                                                   |
+| `normalize.active`  | `bool`                  | `True`        | Whether to use normalization                                                                                                                                            |
+| `normalize.params`  | `dict`                  | `{}`          | Parameters for normalization, see [Normalize](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize) |
+| `augmentations`     | `list[dict]`            | `[]`          | List of `Albumentations` augmentations                                                                                                                                  |
 
 #### Augmentations
 
@@ -306,7 +306,7 @@ trainer:
     # using YAML capture to reuse the image size
     train_image_size: [&height 384, &width 384]
     keep_aspect_ratio: true
-    train_rgb: true
+    color_space: "RGB"
     normalize:
       active: true
     augmentations:
@@ -418,7 +418,7 @@ Each training strategy is a dictionary with the following fields:
 ```yaml
 training_strategy:
   name: "TripleLRSGDStrategy"
-  params: 
+  params:
     warmup_epochs: 3
     warmup_bias_lr: 0.1
     warmup_momentum: 0.8

configs/complex_model.yaml

@@ -100,7 +100,6 @@ trainer:
   preprocessing:
     train_image_size: [&height 384, &width 384]
     keep_aspect_ratio: true
-    train_rgb: true
     normalize:
       active: true
     augmentations:

configs/instance_segmentation_heavy_model.yaml

@@ -0,0 +1,51 @@
+# Example configuration for training a predefined heavy instance segmentation model
+
+model:
+  name: instance_segmentation_heavy
+  predefined_model:
+    name: InstanceSegmentationModel
+    params:
+      variant: heavy
+      loss_params:
+        bbox_loss_weight: 60 # Should be 7.5 * accumulate_grad_batches for best results
+        class_loss_weight: 4 # Should be 0.5 * accumulate_grad_batches for best results
+        dfl_loss_weight: 12 # Should be 1.5 * accumulate_grad_batches for best results
+
+loader:
+  params:
+    dataset_name: coco_test
+
+trainer:
+  preprocessing:
+    train_image_size: [384, 512]
+    keep_aspect_ratio: true
+    normalize:
+      active: true
+
+  batch_size: 8
+  epochs: &epochs 300
+  accumulate_grad_batches: 8 # For best results, always accumulate gradients to effectively use 64 batch size
+  n_workers: 8
+  validation_interval: 10
+  n_log_images: 8
+
+  callbacks:
+    - name: EMACallback
+      params:
+        decay: 0.9999 
+        use_dynamic_decay: True 
+        decay_tau: 2000
+    - name: ExportOnTrainEnd
+    - name: TestOnTrainEnd
+
+  training_strategy:
+    name: "TripleLRSGDStrategy"
+    params: 
+      warmup_epochs: 3
+      warmup_bias_lr: 0.1
+      warmup_momentum: 0.8
+      lr: 0.01
+      lre: 0.0001
+      momentum: 0.937     
+      weight_decay: 0.0005
+      nesterov: True

configs/instance_segmentation_light_model.yaml

@@ -0,0 +1,51 @@
+# Example configuration for training a predefined light instance segmentation model
+
+model:
+  name: instance_segmentation_light
+  predefined_model:
+    name: InstanceSegmentationModel
+    params:
+      variant: light
+      loss_params:
+        bbox_loss_weight: 60 # Should be 7.5 * accumulate_grad_batches for best results
+        class_loss_weight: 4 # Should be 0.5 * accumulate_grad_batches for best results
+        dfl_loss_weight: 12 # Should be 1.5 * accumulate_grad_batches for best results
+
+loader:
+  params:
+    dataset_name: coco_test
+
+trainer:
+  preprocessing:
+    train_image_size: [384, 512]
+    keep_aspect_ratio: true
+    normalize:
+      active: true
+
+  batch_size: 8
+  epochs: &epochs 300
+  accumulate_grad_batches: 8 # For best results, always accumulate gradients to effectively use 64 batch size
+  n_workers: 8
+  validation_interval: 10
+  n_log_images: 8
+
+  callbacks:
+    - name: EMACallback
+      params:
+        decay: 0.9999 
+        use_dynamic_decay: True 
+        decay_tau: 2000
+    - name: ExportOnTrainEnd
+    - name: TestOnTrainEnd
+
+  training_strategy:
+    name: "TripleLRSGDStrategy"
+    params: 
+      warmup_epochs: 3
+      warmup_bias_lr: 0.1
+      warmup_momentum: 0.8
+      lr: 0.01
+      lre: 0.0001
+      momentum: 0.937     
+      weight_decay: 0.0005
+      nesterov: True

luxonis_train/__main__.py

@@ -3,11 +3,10 @@
 from pathlib import Path
 from typing import Annotated
 
+import numpy as np
 import typer
 from luxonis_ml.utils import setup_logging
 
-from luxonis_train.config import Config
-
 setup_logging(use_rich=True)
 
 
@@ -175,47 +174,16 @@ def inspect(
     To close the window press 'q' or 'Esc'.
     """
     import cv2
-    from luxonis_ml.data import Augmentations, LabelType
     from luxonis_ml.data.utils.visualizations import visualize
 
-    from luxonis_train.utils.registry import LOADERS
-
-    cfg = Config.get_config(config, opts)
-    train_augmentations = Augmentations(
-        image_size=cfg.trainer.preprocessing.train_image_size,
-        augmentations=[
-            i.model_dump()
-            for i in cfg.trainer.preprocessing.get_active_augmentations()
-            if i.name != "Normalize"
-        ],
-        train_rgb=cfg.trainer.preprocessing.train_rgb,
-        keep_aspect_ratio=cfg.trainer.preprocessing.keep_aspect_ratio,
-    )
-    val_augmentations = Augmentations(
-        image_size=cfg.trainer.preprocessing.train_image_size,
-        augmentations=[
-            i.model_dump()
-            for i in cfg.trainer.preprocessing.get_active_augmentations()
-        ],
-        train_rgb=cfg.trainer.preprocessing.train_rgb,
-        keep_aspect_ratio=cfg.trainer.preprocessing.keep_aspect_ratio,
-        only_normalize=True,
-    )
+    from luxonis_train.core import LuxonisModel
 
-    Loader = LOADERS.get(cfg.loader.name)
-    loader = Loader(
-        augmentations=(
-            train_augmentations if view == "train" else val_augmentations
-        ),
-        view={
-            "train": cfg.loader.train_view,
-            "val": cfg.loader.val_view,
-            "test": cfg.loader.test_view,
-        }[view],
-        image_source=cfg.loader.image_source,
-        **cfg.loader.params,
-    )
+    opts = opts or []
+    opts.extend(["trainer.preprocessing.normalize.active", "False"])
+
+    model = LuxonisModel(config, opts)
 
+    loader = model.loaders[view.value]
     for images, labels in loader:
         for img in images.values():
             if len(img.shape) != 3:
@@ -226,11 +194,10 @@ def inspect(
             k: v.numpy().transpose(1, 2, 0) for k, v in images.items()
         }
         main_image = np_images[loader.image_source]
-        main_image = cv2.cvtColor(main_image, cv2.COLOR_RGB2BGR)
-        np_labels = {
-            task: (label.numpy(), LabelType(task_type))
-            for task, (label, task_type) in labels.items()
-        }
+        main_image = cv2.cvtColor(main_image, cv2.COLOR_RGB2BGR).astype(
+            np.uint8
+        )
+        np_labels = {task: label.numpy() for task, label in labels.items()}
 
         h, w, _ = main_image.shape
         new_h, new_w = int(h * size_multiplier), int(w * size_multiplier)

luxonis_train/assigners/tal_assigner.py

@@ -143,7 +143,7 @@ def _get_alignment_metric(
         pred_bboxes: Tensor,
         gt_labels: Tensor,
         gt_bboxes: Tensor,
-    ):
+    ) -> tuple[Tensor, Tensor]:
         """Calculates anchor alignment metric and IoU between GTs and
         predicted bboxes.
 
@@ -155,7 +155,11 @@ def _get_alignment_metric(
         @param gt_labels: Initial GT labels [bs, n_max_boxes, 1]
         @type gt_bboxes: Tensor
         @param gt_bboxes: Initial GT bboxes [bs, n_max_boxes, 4]
+        @rtype: tuple[Tensor, Tensor]
+        @return: Anchor alignment metric and IoU between GTs and
+            predicted bboxes.
         """
+
         pred_scores = pred_scores.permute(0, 2, 1)
         gt_labels = gt_labels.to(torch.long)
         ind = torch.zeros([2, self.bs, self.n_max_boxes], dtype=torch.long)
@@ -175,7 +179,7 @@ def _select_topk_candidates(
         metrics: Tensor,
         largest: bool = True,
         topk_mask: Tensor | None = None,
-    ):
+    ) -> Tensor:
         """Selects k anchors based on provided metrics tensor.
 
         @type metrics: Tensor
@@ -250,10 +254,4 @@ def _get_final_assignments(
             torch.full_like(assigned_scores, 0),
         )
 
-        assigned_labels = torch.where(
-            mask_pos_sum.bool(),
-            assigned_labels,
-            torch.full_like(assigned_labels, self.n_classes),
-        )
-
         return assigned_labels, assigned_bboxes, assigned_scores

luxonis_train/attached_modules/base_attached_module.py

@@ -159,18 +159,18 @@ def class_names(self) -> list[str]:
         return self.node.class_names
 
     @property
-    def node_tasks(self) -> dict[TaskType, str]:
+    def node_tasks(self) -> list[TaskType]:
         """Getter for the tasks of the attached node.
 
         @type: dict[TaskType, str]
         @raises RuntimeError: If the node does not have the C{tasks}
             attribute set.
         """
-        if self.node._tasks is None:
+        if self.node.tasks is None:
             raise RuntimeError(
                 "Node must have the `tasks` attribute specified."
             )
-        return self.node._tasks
+        return self.node.tasks
 
     def get_label(
         self, labels: Labels, task_type: TaskType | None = None
@@ -210,13 +210,18 @@ def _get_label(
             if len(self.required_labels) == 1:
                 task_type = self.required_labels[0]
 
-        if task_type is not None:
-            task_name = self.node.get_task_name(task_type)
-            if task_name not in labels:
-                raise IncompatibleException.from_missing_task(
-                    task_type.value, list(labels.keys()), self.name
+        if task_type is not None and self.node.task_name is not None:
+            task_name = self.node.task_name
+            task = f"{task_name}/{task_type.value}"
+            if task not in labels:
+                raise IncompatibleException(
+                    f"Module '{self.name}' requires label of type "
+                    f"'{task_type.value}' assigned to task '{task_name}', "
+                    "but the label is missing from the dataset. "
+                    f"Available labels: {list(labels.keys())}. "
+                    f"Missing label: '{task}'."
                 )
-            return labels[task_name]
+            return labels[task], task_type
 
         raise ValueError(
             f"{self.name} requires multiple labels. You must provide the "
@@ -260,7 +265,7 @@ def get_input_tensors(
                         f"Task {task_type.value} is not supported by the node "
                         f"{self.node.name}."
                     )
-                return inputs[self.node_tasks[task_type]]
+                return inputs[f"{self.node.task_name}/{task_type.value}"]
             else:
                 if task_type not in inputs:
                     raise IncompatibleException(
@@ -273,7 +278,20 @@ def get_input_tensors(
                 f"{self.name} requires multiple labels, "
                 "you must provide the `task_type` argument to extract the desired input."
             )
-        return inputs[self.node_tasks[self.required_labels[0]]]
+        if len(self.node_tasks) == 1:
+            task_type = self.node_tasks[0].value
+        else:
+            required_label = self.required_labels[0]
+            for task in self.node_tasks:
+                if task.value == required_label:
+                    task_type = task.value
+                    break
+            else:
+                raise IncompatibleException(
+                    f"Task {required_label} is not supported by the node "
+                    f"{self.node.name}."
+                )
+        return inputs[f"{self.node.task_name}/{task_type}"]
 
     def prepare(
         self, inputs: Packet[Tensor], labels: Labels | None
@@ -305,7 +323,7 @@ def prepare(
         @raises RuntimeError: If the C{tasks} attribute is not set on the node.
         @raises RuntimeError: If the C{supported_tasks} attribute is not set on the module.
         """
-        if self.node._tasks is None:
+        if self.node.tasks is None:
             raise RuntimeError(
                 f"{self.node.name} must have the `tasks` attribute specified "
                 f"for {self.name} to make use of the default `prepare` method."