175v2 add consep segmentation dataset

.gitignore

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

configs/vision/dino_vit/online/consep.yaml

@@ -0,0 +1,97 @@
+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/consep/${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224}}
+    max_steps: &MAX_STEPS ${oc.env:MAX_STEPS, 513}
+    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 8
+    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.Consep
+        init_args: &DATASET_ARGS
+          root: ${oc.env:DATA_ROOT, ./data}/consep
+          split: train
+          transforms:
+            class_path: eva.vision.data.transforms.common.ResizeAndClamp
+            init_args:
+              mean: *NORMALIZE_MEAN
+              std: *NORMALIZE_STD
+      val:
+        class_path: eva.vision.datasets.Consep
+        init_args:
+          <<: *DATASET_ARGS
+          split: val
+    dataloaders:
+      train:
+        batch_size: &BATCH_SIZE ${oc.env:BATCH_SIZE, 16}
+        shuffle: true
+      val:
+        batch_size: *BATCH_SIZE

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

@@ -45,6 +45,7 @@ dependencies = [
     "onnx>=1.16.0",
     "toolz>=0.12.1",
     "rich>=13.7.1",
+    "scipy>=1.13.1",
 ]
 
 [project.urls]
@@ -60,15 +61,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,
+    )