Support semantic segmentation downstream evaluation tasks

.gitignore

@@ -173,3 +173,9 @@ cython_debug/
 
 # ignore local data
 /data/
+
+# numpy data
+*.npy
+
+# NiFti data
+*.nii.gz

configs/vision/dino_vit/online/total_segmentator_2d.yaml

@@ -0,0 +1,109 @@
+trainer:
+  class_path: eva.Trainer
+  init_args:
+    n_runs: &N_RUNS ${oc.env:N_RUNS, 1}
+    default_root_dir: &OUTPUT_ROOT ${oc.env:OUTPUT_ROOT, logs/total_segmentator_2d/${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224}}
+    max_steps: &MAX_STEPS ${oc.env:MAX_STEPS, 20000}
+    callbacks:
+      - class_path: eva.vision.callbacks.SemanticSegmentationLogger
+        init_args:
+          log_every_n_epochs: 1
+          mean: &NORMALIZE_MEAN [0.5, 0.5, 0.5]
+          std: &NORMALIZE_STD [0.5, 0.5, 0.5]
+      - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+        init_args:
+          filename: best
+          save_last: true
+          save_top_k: 1
+          monitor: &MONITOR_METRIC ${oc.env:MONITOR_METRIC, val/MulticlassJaccardIndex}
+          mode: &MONITOR_METRIC_MODE ${oc.env:MONITOR_METRIC_MODE, max}
+      - class_path: lightning.pytorch.callbacks.EarlyStopping
+        init_args:
+          min_delta: 0
+          patience: 5
+          monitor: *MONITOR_METRIC
+          mode: *MONITOR_METRIC_MODE
+    logger:
+      - class_path: lightning.pytorch.loggers.TensorBoardLogger
+        init_args:
+          save_dir: *OUTPUT_ROOT
+          name: ""
+model:
+  class_path: eva.vision.models.modules.SemanticSegmentationModule
+  init_args:
+    encoder:
+      class_path: eva.vision.models.networks.encoders.TimmEncoder
+      init_args:
+        model_name: ${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224}
+        pretrained: true
+        out_indices: ${oc.env:TIMM_MODEL_OUT_INDICES, 1}
+        model_arguments:
+          dynamic_img_size: true
+    decoder:
+      class_path: eva.vision.models.networks.decoders.segmentation.ConvDecoderMS
+      init_args:
+        in_features: ${oc.env:DECODER_IN_FEATURES, 384}
+        num_classes:  &NUM_CLASSES 118
+    criterion: torch.nn.CrossEntropyLoss
+    lr_multiplier_encoder: 0.0
+    optimizer:
+      class_path: torch.optim.AdamW
+      init_args:
+        lr: 0.0001
+        weight_decay: 0.05
+    lr_scheduler:
+      class_path: torch.optim.lr_scheduler.PolynomialLR
+      init_args:
+        total_iters: *MAX_STEPS
+        power: 0.9
+    metrics:
+      common:
+        - class_path: eva.metrics.AverageLoss
+      evaluation:
+        - class_path: eva.core.metrics.defaults.MulticlassSegmentationMetrics
+          init_args:
+            num_classes: *NUM_CLASSES
+        - class_path: eva.core.metrics.wrappers.ClasswiseWrapper
+          init_args:
+            metric:
+              class_path: torchmetrics.classification.MulticlassF1Score
+              init_args:
+                num_classes: *NUM_CLASSES
+                average: null
+data:
+  class_path: eva.DataModule
+  init_args:
+    datasets:
+      train:
+        class_path: eva.vision.datasets.TotalSegmentator2D
+        init_args: &DATASET_ARGS
+          root: ${oc.env:DATA_ROOT, ./data}/total_segmentator
+          split: train
+          download: ${oc.env:DOWNLOAD_DATA, false}
+          # Set `download: true` to download the dataset from https://zenodo.org/records/10047292
+          # The TotalSegmentator dataset is distributed under the following license: 
+          # "Creative Commons Attribution 4.0 International"
+          # (see: https://creativecommons.org/licenses/by/4.0/deed.en)
+          transforms:
+            class_path: eva.vision.data.transforms.common.ResizeAndClamp
+            init_args:
+              mean: *NORMALIZE_MEAN
+              std: *NORMALIZE_STD
+      val:
+        class_path: eva.vision.datasets.TotalSegmentator2D
+        init_args:
+          <<: *DATASET_ARGS
+          split: val
+      test:
+        class_path: eva.vision.datasets.TotalSegmentator2D
+        init_args:
+          <<: *DATASET_ARGS
+          split: test
+    dataloaders:
+      train:
+        batch_size: &BATCH_SIZE ${oc.env:BATCH_SIZE, 16}
+        shuffle: true
+      val:
+        batch_size: *BATCH_SIZE
+      test:
+        batch_size: *BATCH_SIZE

pyproject.toml

@@ -60,15 +60,15 @@ vision = [
     "h5py>=3.10.0",
     "nibabel>=5.2.0",
     "opencv-python-headless>=4.9.0.80",
-    "timm>=0.9.12",
+    "timm @ git+https://github.com/huggingface/pytorch-image-models.git@main",
     "torchvision>=0.17.0",
     "openslide-python>=1.3.1",
 ]
 all = [
     "h5py>=3.10.0",
     "nibabel>=5.2.0",
     "opencv-python-headless>=4.9.0.80",
-    "timm>=0.9.12",
+    "timm @ git+https://github.com/huggingface/pytorch-image-models.git@main",
     "torchvision>=0.17.0",
     "openslide-python>=1.3.1",
 ]

src/eva/core/callbacks/writers/__init__.py

@@ -1,4 +1,4 @@
-"""Callbacks API."""
+"""Writers callbacks API."""
 
 from eva.core.callbacks.writers.embeddings import ClassificationEmbeddingsWriter
 

src/eva/core/loggers/log/__init__.py

@@ -1,5 +1,6 @@
-"""Experiment loggers actions."""
+"""Experiment loggers operations."""
 
+from eva.core.loggers.log.image import log_image
 from eva.core.loggers.log.parameters import log_parameters
 
-__all__ = ["log_parameters"]
+__all__ = ["log_image", "log_parameters"]

src/eva/core/loggers/log/image.py

@@ -0,0 +1,59 @@
+"""Image log functionality."""
+
+import functools
+
+import torch
+
+from eva.core.loggers import loggers
+from eva.core.loggers.log import utils
+
+
+@functools.singledispatch
+def log_image(
+    logger,
+    tag: str,
+    image: torch.Tensor,
+    step: int = 0,
+) -> None:
+    """Adds an image to the logger.
+
+    Args:
+        logger: The desired logger.
+        tag: The log tag.
+        image: The image tensor to log. It should have
+            the shape of (3,H,W) and (0,1) normalized.
+        step: The global step of the log.
+    """
+    utils.raise_not_supported(logger, "image")
+
+
+@log_image.register
+def _(
+    loggers: list,
+    tag: str,
+    image: torch.Tensor,
+    step: int = 0,
+) -> None:
+    """Adds an image to a list of supported loggers."""
+    for logger in loggers:
+        log_image(
+            logger,
+            tag=tag,
+            image=image,
+            step=step,
+        )
+
+
+@log_image.register
+def _(
+    logger: loggers.TensorBoardLogger,
+    tag: str,
+    image: torch.Tensor,
+    step: int = 0,
+) -> None:
+    """Adds an image to a TensorBoard logger."""
+    logger.experiment.add_image(
+        tag=tag,
+        img_tensor=image,
+        global_step=step,
+    )

src/eva/core/loggers/loggers.py

@@ -0,0 +1,6 @@
+"""Experimental loggers."""
+
+from lightning.pytorch.loggers import TensorBoardLogger
+
+Loggers = TensorBoardLogger
+"""Supported loggers."""

src/eva/core/metrics/defaults/__init__.py

@@ -1,6 +1,13 @@
 """Default metric collections API."""
 
-from eva.core.metrics.defaults.classification.binary import BinaryClassificationMetrics
-from eva.core.metrics.defaults.classification.multiclass import MulticlassClassificationMetrics
+from eva.core.metrics.defaults.classification import (
+    BinaryClassificationMetrics,
+    MulticlassClassificationMetrics,
+)
+from eva.core.metrics.defaults.segmentation import MulticlassSegmentationMetrics
 
-__all__ = ["MulticlassClassificationMetrics", "BinaryClassificationMetrics"]
+__all__ = [
+    "MulticlassClassificationMetrics",
+    "BinaryClassificationMetrics",
+    "MulticlassSegmentationMetrics",
+]

src/eva/core/metrics/defaults/classification/__init__.py

@@ -3,4 +3,4 @@
 from eva.core.metrics.defaults.classification.binary import BinaryClassificationMetrics
 from eva.core.metrics.defaults.classification.multiclass import MulticlassClassificationMetrics
 
-__all__ = ["MulticlassClassificationMetrics", "BinaryClassificationMetrics"]
+__all__ = ["BinaryClassificationMetrics", "MulticlassClassificationMetrics"]

src/eva/core/metrics/defaults/classification/binary.py

@@ -17,15 +17,6 @@ def __init__(
     ) -> None:
         """Initializes the binary classification metrics.
 
-        The metrics instantiated here are:
-
-        - BinaryAUROC
-        - BinaryAccuracy
-        - BinaryBalancedAccuracy
-        - BinaryF1Score
-        - BinaryPrecision
-        - BinaryRecall
-
         Args:
             threshold: Threshold for transforming probability to binary (0,1) predictions
             ignore_index: Specifies a target value that is ignored and does not

src/eva/core/metrics/defaults/classification/multiclass.py

@@ -20,14 +20,6 @@ def __init__(
     ) -> None:
         """Initializes the multi-class classification metrics.
 
-        The metrics instantiated here are:
-
-        - MulticlassAccuracy
-        - MulticlassPrecision
-        - MulticlassRecall
-        - MulticlassF1Score
-        - MulticlassAUROC
-
         Args:
             num_classes: Integer specifying the number of classes.
             average: Defines the reduction that is applied over labels.

src/eva/core/metrics/defaults/segmentation/__init__.py

@@ -0,0 +1,5 @@
+"""Default segmentation metric collections API."""
+
+from eva.core.metrics.defaults.segmentation.multiclass import MulticlassSegmentationMetrics
+
+__all__ = ["MulticlassSegmentationMetrics"]

src/eva/core/metrics/defaults/segmentation/multiclass.py

@@ -0,0 +1,66 @@
+"""Default metric collection for multiclass semantic segmentation tasks."""
+
+from typing import Literal
+
+from torchmetrics import classification
+
+from eva.core.metrics import structs
+
+
+class MulticlassSegmentationMetrics(structs.MetricCollection):
+    """Default metrics for multi-class semantic segmentation tasks."""
+
+    def __init__(
+        self,
+        num_classes: int,
+        average: Literal["macro", "weighted", "none"] = "macro",
+        ignore_index: int | None = None,
+        prefix: str | None = None,
+        postfix: str | None = None,
+    ) -> None:
+        """Initializes the multi-class semantic segmentation metrics.
+
+        Args:
+            num_classes: Integer specifying the number of classes.
+            average: Defines the reduction that is applied over labels.
+            ignore_index: Specifies a target value that is ignored and
+                does not contribute to the metric calculation.
+            prefix: A string to add before the keys in the output dictionary.
+            postfix: A string to add after the keys in the output dictionary.
+        """
+        super().__init__(
+            metrics=[
+                classification.MulticlassJaccardIndex(
+                    num_classes=num_classes,
+                    average=average,
+                    ignore_index=ignore_index,
+                ),
+                classification.MulticlassF1Score(
+                    num_classes=num_classes,
+                    average=average,
+                    ignore_index=ignore_index,
+                ),
+                classification.MulticlassPrecision(
+                    num_classes=num_classes,
+                    average=average,
+                    ignore_index=ignore_index,
+                ),
+                classification.MulticlassRecall(
+                    num_classes=num_classes,
+                    average=average,
+                    ignore_index=ignore_index,
+                ),
+            ],
+            prefix=prefix,
+            postfix=postfix,
+            compute_groups=[
+                [
+                    "MulticlassJaccardIndex",
+                ],
+                [
+                    "MulticlassF1Score",
+                    "MulticlassPrecision",
+                    "MulticlassRecall",
+                ],
+            ],
+        )

src/eva/core/metrics/structs/schemas.py

@@ -44,4 +44,6 @@ def _join_with_common(self, metrics: MetricModuleType | None) -> MetricModuleTyp
         if metrics is None or self.common is None:
             return self.common or metrics
 
-        return [self.common, metrics]  # type: ignore
+        metrics = metrics if isinstance(metrics, list) else [metrics]  # type: ignore
+        common = self.common if isinstance(self.common, list) else [self.common]
+        return common + metrics  # type: ignore

src/eva/core/metrics/wrappers/__init__.py

@@ -0,0 +1,5 @@
+"""Metric wrappers API."""
+
+from eva.core.metrics.wrappers.classwise import ClasswiseWrapper
+
+__all__ = ["ClasswiseWrapper"]