add confusion matrix to predefined models

luxonis_train/attached_modules/metrics/confusion_matrix.py

@@ -2,7 +2,10 @@
 
 import torch
 from torch import Tensor
-from torchmetrics.classification import MulticlassConfusionMatrix
+from torchmetrics.classification import (
+    BinaryConfusionMatrix,
+    MulticlassConfusionMatrix,
+)
 from torchvision.ops import box_convert, box_iou
 
 from luxonis_train.enums import TaskType
@@ -71,9 +74,12 @@ def __init__(
 
         self.metric_cm = None
         if self.is_classification or self.is_segmentation:
-            self.metric_cm = MulticlassConfusionMatrix(
-                num_classes=self.n_classes
-            )
+            if self.n_classes == 1:
+                self.metric_cm = BinaryConfusionMatrix()
+            else:
+                self.metric_cm = MulticlassConfusionMatrix(
+                    num_classes=self.n_classes
+                )
 
         if self.is_detection:
             self.add_state(
@@ -153,16 +159,32 @@ def update(
         if "classification" in predictions and "classification" in targets:
             preds = predictions["classification"]
             target = targets["classification"]
-            pred_classes = preds[0].argmax(dim=1)  # [B]
-            target_classes = target.argmax(dim=1)  # [B]
+            pred_classes = (
+                preds[0].argmax(dim=1)
+                if preds[0].shape[1] > 1
+                else preds[0].sigmoid().squeeze(1).round().int()
+            )  # [B]
+            target_classes = (
+                target.argmax(dim=1)
+                if target.shape[1] > 1
+                else target.squeeze(1).round().int()
+            )  # [B]
             if self.metric_cm is not None:
                 self.metric_cm.update(pred_classes, target_classes)
 
         if "segmentation" in predictions and "segmentation" in targets:
             preds = predictions["segmentation"]
             target = targets["segmentation"]
-            pred_masks = preds[0].argmax(dim=1)  # [B, H, W]
-            target_masks = target.argmax(dim=1)  # [B, H, W]
+            pred_masks = (
+                preds[0].argmax(dim=1)
+                if preds[0].shape[1] > 1
+                else preds[0].squeeze(1).sigmoid().round().int()
+            )  # [B, H, W]
+            target_masks = (
+                target.argmax(dim=1)
+                if target.shape[1] > 1
+                else target.squeeze(1).round().int()
+            )  # [B, H, W]
             if self.metric_cm is not None:
                 self.metric_cm.update(
                     pred_masks.view(-1), target_masks.view(-1)

luxonis_train/config/predefined_models/classification_model.py

@@ -53,6 +53,8 @@ def __init__(
         visualizer_params: Params | None = None,
         task: Literal["multiclass", "multilabel"] = "multiclass",
         task_name: str | None = None,
+        enable_confusion_matrix: bool = True,
+        confusion_matrix_params: Params | None = None,
     ):
         var_config = get_variant(variant)
 
@@ -67,6 +69,8 @@ def __init__(
         self.visualizer_params = visualizer_params or {}
         self.task = task
         self.task_name = task_name or "classification"
+        self.enable_confusion_matrix = enable_confusion_matrix
+        self.confusion_matrix_params = confusion_matrix_params or {}
 
     @property
     def nodes(self) -> list[ModelNodeConfig]:
@@ -104,7 +108,7 @@ def losses(self) -> list[LossModuleConfig]:
     @property
     def metrics(self) -> list[MetricModuleConfig]:
         """Defines the metrics used for evaluation."""
-        return [
+        metrics = [
             MetricModuleConfig(
                 name="F1Score",
                 alias=f"F1Score-{self.task_name}",
@@ -125,6 +129,16 @@ def metrics(self) -> list[MetricModuleConfig]:
                 params={"task": self.task},
             ),
         ]
+        if self.enable_confusion_matrix:
+            metrics.append(
+                MetricModuleConfig(
+                    name="ConfusionMatrix",
+                    alias=f"ConfusionMatrix-{self.task_name}",
+                    attached_to=f"ClassificationHead-{self.task_name}",
+                    params={**self.confusion_matrix_params},
+                )
+            )
+        return metrics
 
     @property
     def visualizers(self) -> list[AttachedModuleConfig]:

luxonis_train/config/predefined_models/detection_model.py

@@ -66,6 +66,8 @@ def __init__(
         loss_params: Params | None = None,
         visualizer_params: Params | None = None,
         task_name: str | None = None,
+        enable_confusion_matrix: bool = True,
+        confusion_matrix_params: Params | None = None,
     ):
         var_config = get_variant(variant)
 
@@ -81,6 +83,8 @@ def __init__(
         self.loss_params = loss_params or {"n_warmup_epochs": 0}
         self.visualizer_params = visualizer_params or {}
         self.task_name = task_name or "boundingbox"
+        self.enable_confusion_matrix = enable_confusion_matrix
+        self.confusion_matrix_params = confusion_matrix_params or {}
 
     @property
     def nodes(self) -> list[ModelNodeConfig]:
@@ -135,14 +139,24 @@ def losses(self) -> list[LossModuleConfig]:
     @property
     def metrics(self) -> list[MetricModuleConfig]:
         """Defines the metrics used for evaluation."""
-        return [
+        metrics = [
             MetricModuleConfig(
                 name="MeanAveragePrecision",
                 alias=f"MeanAveragePrecision-{self.task_name}",
                 attached_to=f"EfficientBBoxHead-{self.task_name}",
                 is_main_metric=True,
             ),
         ]
+        if self.enable_confusion_matrix:
+            metrics.append(
+                MetricModuleConfig(
+                    name="ConfusionMatrix",
+                    alias=f"ConfusionMatrix-{self.task_name}",
+                    attached_to=f"EfficientBBoxHead-{self.task_name}",
+                    params={**self.confusion_matrix_params},
+                )
+            )
+        return metrics
 
     @property
     def visualizers(self) -> list[AttachedModuleConfig]:

luxonis_train/config/predefined_models/keypoint_detection_model.py

@@ -64,6 +64,8 @@ def __init__(
         bbox_visualizer_params: Params | None = None,
         bbox_task_name: str | None = None,
         kpt_task_name: str | None = None,
+        enable_confusion_matrix: bool = True,
+        confusion_matrix_params: Params | None = None,
     ):
         var_config = get_variant(variant)
 
@@ -81,6 +83,8 @@ def __init__(
         self.bbox_visualizer_params = bbox_visualizer_params or {}
         self.bbox_task_name = bbox_task_name or "boundingbox"
         self.kpt_task_name = kpt_task_name or "keypoints"
+        self.enable_confusion_matrix = enable_confusion_matrix
+        self.confusion_matrix_params = confusion_matrix_params or {}
 
     @property
     def nodes(self) -> list[ModelNodeConfig]:
@@ -143,7 +147,7 @@ def losses(self) -> list[LossModuleConfig]:
     @property
     def metrics(self) -> list[MetricModuleConfig]:
         """Defines the metrics used for evaluation."""
-        return [
+        metrics = [
             MetricModuleConfig(
                 name="ObjectKeypointSimilarity",
                 alias=f"ObjectKeypointSimilarity-{self.kpt_task_name}",
@@ -156,6 +160,16 @@ def metrics(self) -> list[MetricModuleConfig]:
                 attached_to=f"EfficientKeypointBBoxHead-{self.kpt_task_name}",
             ),
         ]
+        if self.enable_confusion_matrix:
+            metrics.append(
+                MetricModuleConfig(
+                    name="ConfusionMatrix",
+                    alias=f"ConfusionMatrix-{self.kpt_task_name}",
+                    attached_to=f"EfficientKeypointBBoxHead-{self.kpt_task_name}",
+                    params={**self.confusion_matrix_params},
+                )
+            )
+        return metrics
 
     @property
     def visualizers(self) -> list[AttachedModuleConfig]:

luxonis_train/config/predefined_models/segmentation_model.py

@@ -57,6 +57,8 @@ def __init__(
         visualizer_params: Params | None = None,
         task: Literal["binary", "multiclass"] = "binary",
         task_name: str | None = None,
+        enable_confusion_matrix: bool = True,
+        confusion_matrix_params: Params | None = None,
     ):
         var_config = get_variant(variant)
 
@@ -72,6 +74,8 @@ def __init__(
         self.visualizer_params = visualizer_params or {}
         self.task = task
         self.task_name = task_name or "segmentation"
+        self.enable_confusion_matrix = enable_confusion_matrix
+        self.confusion_matrix_params = confusion_matrix_params or {}
 
     @property
     def nodes(self) -> list[ModelNodeConfig]:
@@ -154,7 +158,7 @@ def losses(self) -> list[LossModuleConfig]:
     @property
     def metrics(self) -> list[MetricModuleConfig]:
         """Defines the metrics used for evaluation."""
-        return [
+        metrics = [
             MetricModuleConfig(
                 name="JaccardIndex",
                 alias=f"JaccardIndex-{self.task_name}",
@@ -169,6 +173,16 @@ def metrics(self) -> list[MetricModuleConfig]:
                 params={"task": self.task},
             ),
         ]
+        if self.enable_confusion_matrix:
+            metrics.append(
+                MetricModuleConfig(
+                    name="ConfusionMatrix",
+                    alias=f"ConfusionMatrix-{self.task_name}",
+                    attached_to=f"DDRNetSegmentationHead-{self.task_name}",
+                    params={**self.confusion_matrix_params},
+                )
+            )
+        return metrics
 
     @property
     def visualizers(self) -> list[AttachedModuleConfig]: