From 567baf4fb17ad30f6986afa6594552cbaddc318e Mon Sep 17 00:00:00 2001 From: Jernej Sabadin Date: Thu, 9 Jan 2025 14:37:17 +0100 Subject: [PATCH 1/3] add confusion matrix to predefined models --- .../predefined_models/classification_model.py | 18 ++++++++++++++++-- .../predefined_models/detection_model.py | 18 ++++++++++++++++-- .../predefined_models/segmentation_model.py | 18 ++++++++++++++++-- 3 files changed, 48 insertions(+), 6 deletions(-) diff --git a/luxonis_train/config/predefined_models/classification_model.py b/luxonis_train/config/predefined_models/classification_model.py index e028bba5..ca4ad62b 100644 --- a/luxonis_train/config/predefined_models/classification_model.py +++ b/luxonis_train/config/predefined_models/classification_model.py @@ -1,4 +1,4 @@ -from typing import Literal, TypeAlias +from typing import Literal, Optional, TypeAlias from pydantic import BaseModel @@ -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: Optional[Params] = 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]: diff --git a/luxonis_train/config/predefined_models/detection_model.py b/luxonis_train/config/predefined_models/detection_model.py index dbbc8886..fbd6a853 100644 --- a/luxonis_train/config/predefined_models/detection_model.py +++ b/luxonis_train/config/predefined_models/detection_model.py @@ -1,4 +1,4 @@ -from typing import Literal, TypeAlias +from typing import Literal, Optional, TypeAlias from pydantic import BaseModel @@ -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: Optional[Params] = 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,7 +139,7 @@ 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}", @@ -143,6 +147,16 @@ def metrics(self) -> list[MetricModuleConfig]: 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]: diff --git a/luxonis_train/config/predefined_models/segmentation_model.py b/luxonis_train/config/predefined_models/segmentation_model.py index eff4fd02..815b29a9 100644 --- a/luxonis_train/config/predefined_models/segmentation_model.py +++ b/luxonis_train/config/predefined_models/segmentation_model.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Literal, Optional from pydantic import BaseModel @@ -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: Optional[Params] = 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]: From 44e11eb378000175618705b689cff6c78794aea9 Mon Sep 17 00:00:00 2001 From: Jernej Sabadin Date: Thu, 9 Jan 2025 14:44:21 +0100 Subject: [PATCH 2/3] add CM metric to kpts predefined model --- .../predefined_models/classification_model.py | 4 ++-- .../config/predefined_models/detection_model.py | 4 ++-- .../keypoint_detection_model.py | 16 +++++++++++++++- .../predefined_models/segmentation_model.py | 4 ++-- 4 files changed, 21 insertions(+), 7 deletions(-) diff --git a/luxonis_train/config/predefined_models/classification_model.py b/luxonis_train/config/predefined_models/classification_model.py index ca4ad62b..56cc13d3 100644 --- a/luxonis_train/config/predefined_models/classification_model.py +++ b/luxonis_train/config/predefined_models/classification_model.py @@ -1,4 +1,4 @@ -from typing import Literal, Optional, TypeAlias +from typing import Literal, TypeAlias from pydantic import BaseModel @@ -54,7 +54,7 @@ def __init__( task: Literal["multiclass", "multilabel"] = "multiclass", task_name: str | None = None, enable_confusion_matrix: bool = True, - confusion_matrix_params: Optional[Params] = None, + confusion_matrix_params: Params | None = None, ): var_config = get_variant(variant) diff --git a/luxonis_train/config/predefined_models/detection_model.py b/luxonis_train/config/predefined_models/detection_model.py index fbd6a853..82cfd98b 100644 --- a/luxonis_train/config/predefined_models/detection_model.py +++ b/luxonis_train/config/predefined_models/detection_model.py @@ -1,4 +1,4 @@ -from typing import Literal, Optional, TypeAlias +from typing import Literal, TypeAlias from pydantic import BaseModel @@ -67,7 +67,7 @@ def __init__( visualizer_params: Params | None = None, task_name: str | None = None, enable_confusion_matrix: bool = True, - confusion_matrix_params: Optional[Params] = None, + confusion_matrix_params: Params | None = None, ): var_config = get_variant(variant) diff --git a/luxonis_train/config/predefined_models/keypoint_detection_model.py b/luxonis_train/config/predefined_models/keypoint_detection_model.py index 51d790a7..1cdcfa05 100644 --- a/luxonis_train/config/predefined_models/keypoint_detection_model.py +++ b/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]: diff --git a/luxonis_train/config/predefined_models/segmentation_model.py b/luxonis_train/config/predefined_models/segmentation_model.py index 815b29a9..554ead9e 100644 --- a/luxonis_train/config/predefined_models/segmentation_model.py +++ b/luxonis_train/config/predefined_models/segmentation_model.py @@ -1,4 +1,4 @@ -from typing import Literal, Optional +from typing import Literal from pydantic import BaseModel @@ -58,7 +58,7 @@ def __init__( task: Literal["binary", "multiclass"] = "binary", task_name: str | None = None, enable_confusion_matrix: bool = True, - confusion_matrix_params: Optional[Params] = None, + confusion_matrix_params: Params | None = None, ): var_config = get_variant(variant) From 218fcc5fc58752da10a1b6e40fa95a2059534811 Mon Sep 17 00:00:00 2001 From: Jernej Sabadin Date: Thu, 9 Jan 2025 18:18:50 +0100 Subject: [PATCH 3/3] support for n_classes = 1 --- .../metrics/confusion_matrix.py | 38 +++++++++++++++---- 1 file changed, 30 insertions(+), 8 deletions(-) diff --git a/luxonis_train/attached_modules/metrics/confusion_matrix.py b/luxonis_train/attached_modules/metrics/confusion_matrix.py index 079a9315..f719e084 100644 --- a/luxonis_train/attached_modules/metrics/confusion_matrix.py +++ b/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)