Add support for offline total segmentator training

configs/vision/dino_vit/offline/consep.yaml

@@ -0,0 +1,124 @@
+---
+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/${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224_dino}/consep}
+    max_steps: &MAX_STEPS ${oc.env:MAX_STEPS, 513}
+    log_every_n_steps: 6
+    callbacks:
+      - class_path: eva.callbacks.ConfigurationLogger
+      - class_path: eva.vision.callbacks.SemanticSegmentationLogger
+        init_args:
+          log_every_n_epochs: 1
+          log_images: false
+      - 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: 100
+          monitor: *MONITOR_METRIC
+          mode: *MONITOR_METRIC_MODE
+      - class_path: eva.callbacks.SegmentationEmbeddingsWriter
+        init_args:
+          output_dir: &DATASET_EMBEDDINGS_ROOT ${oc.env:EMBEDDINGS_ROOT, ./data/embeddings}/${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224_dino}/consep
+          dataloader_idx_map:
+            0: train
+            1: val
+          metadata_keys: ["coords"]
+          backbone:
+            class_path: eva.vision.models.networks.encoders.TimmEncoder
+            init_args:
+              model_name: ${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224_dino}
+              pretrained: ${oc.env:MODEL_PRETRAINED, true}
+              out_indices: ${oc.env:TIMM_MODEL_OUT_INDICES, 3}
+              checkpoint_path: &CHECKPOINT_PATH ${oc.env:CHECKPOINT_PATH, null}
+              model_arguments:
+                dynamic_img_size: true
+    logger:
+      - class_path: lightning.pytorch.loggers.TensorBoardLogger
+        init_args:
+          save_dir: *OUTPUT_ROOT
+          name: ""
+model:
+  class_path: eva.vision.models.modules.SemanticSegmentationModule
+  init_args:
+    decoder:
+      class_path: eva.vision.models.networks.decoders.segmentation.ConvDecoderMS
+      init_args:
+        in_features: ${oc.env:DECODER_IN_FEATURES, 1152}
+        num_classes:  &NUM_CLASSES 8
+    criterion: torch.nn.CrossEntropyLoss
+    lr_multiplier_encoder: 0.0
+    optimizer:
+      class_path: torch.optim.AdamW
+      init_args:
+        lr: 0.001
+        weight_decay: 0.01
+    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.EmbeddingsSegmentationDataset
+        init_args: &DATASET_ARGS
+          root: *DATASET_EMBEDDINGS_ROOT
+          manifest_file: manifest.csv
+          split: train
+      val:
+        class_path: eva.vision.datasets.EmbeddingsSegmentationDataset
+        init_args:
+          <<: *DATASET_ARGS
+          split: val
+      predict:
+        - class_path: eva.vision.datasets.CoNSeP
+          init_args: &PREDICT_DATASET_ARGS
+            root: ${oc.env:DATA_ROOT, ./data}/consep
+            split: train
+            sampler:
+              class_path: eva.vision.data.wsi.patching.samplers.ForegroundGridSampler
+              init_args:
+                max_samples: 25
+            transforms:
+              class_path: eva.vision.data.transforms.common.ResizeAndCrop
+              init_args:
+                size: ${oc.env:RESIZE_DIM, 224}  
+                mean: &NORMALIZE_MEAN ${oc.env:NORMALIZE_MEAN, [0.485, 0.456, 0.406]} 
+                std: &NORMALIZE_STD ${oc.env:NORMALIZE_STD, [0.229, 0.224, 0.225]}
+        - class_path: eva.vision.datasets.CoNSeP
+          init_args:
+            <<: *PREDICT_DATASET_ARGS
+            split: val
+    dataloaders:
+      train:
+        batch_size: &BATCH_SIZE ${oc.env:BATCH_SIZE, 64}
+        shuffle: true
+      val:
+        batch_size: *BATCH_SIZE
+      predict:
+        batch_size: &PREDICT_BATCH_SIZE ${oc.env:PREDICT_BATCH_SIZE, 64}

configs/vision/dino_vit/offline/monusac.yaml

@@ -0,0 +1,124 @@
+---
+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/${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224_dino}/monusac}
+    max_steps: &MAX_STEPS ${oc.env:MAX_STEPS, 513}
+    log_every_n_steps: 6
+    callbacks:
+      - class_path: eva.callbacks.ConfigurationLogger
+      - class_path: eva.vision.callbacks.SemanticSegmentationLogger
+        init_args:
+          log_every_n_epochs: 1
+          log_images: false
+      - 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: 100
+          monitor: *MONITOR_METRIC
+          mode: *MONITOR_METRIC_MODE
+      - class_path: eva.callbacks.SegmentationEmbeddingsWriter
+        init_args:
+          output_dir: &DATASET_EMBEDDINGS_ROOT ${oc.env:EMBEDDINGS_ROOT, ./data/embeddings}/${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224_dino}/monusac
+          dataloader_idx_map:
+            0: train
+            1: test
+          backbone:
+            class_path: eva.vision.models.networks.encoders.TimmEncoder
+            init_args:
+              model_name: ${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224_dino}
+              pretrained: ${oc.env:MODEL_PRETRAINED, true}
+              out_indices: ${oc.env:TIMM_MODEL_OUT_INDICES, 1}
+              checkpoint_path: &CHECKPOINT_PATH ${oc.env:CHECKPOINT_PATH, null}
+              model_arguments:
+                dynamic_img_size: true
+    logger:
+      - class_path: lightning.pytorch.loggers.TensorBoardLogger
+        init_args:
+          save_dir: *OUTPUT_ROOT
+          name: ""
+model:
+  class_path: eva.vision.models.modules.SemanticSegmentationModule
+  init_args:
+    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.001
+        weight_decay: 0.01
+    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.EmbeddingsSegmentationDataset
+        init_args: &DATASET_ARGS
+          root: *DATASET_EMBEDDINGS_ROOT
+          manifest_file: manifest.csv
+          split: train
+      val:
+        class_path: eva.vision.datasets.EmbeddingsSegmentationDataset
+        init_args:
+          <<: *DATASET_ARGS
+          split: test
+      predict:
+        - class_path: eva.vision.datasets.MoNuSAC
+          init_args: &PREDICT_DATASET_ARGS
+            root: ${oc.env:DATA_ROOT, ./data}/monusac
+            split: train
+            download: ${oc.env:DOWNLOAD_DATA, false}
+            # Set `download: true` to download the dataset from https://monusac-2020.grand-challenge.org/Data/
+            # The MoNuSAC dataset is distributed under the following license: 
+            # "Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International"
+            # (see: https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode)
+            transforms:
+              class_path: eva.vision.data.transforms.common.ResizeAndCrop
+              init_args:
+                size: ${oc.env:RESIZE_DIM, 224}  
+                mean: &NORMALIZE_MEAN ${oc.env:NORMALIZE_MEAN, [0.485, 0.456, 0.406]} 
+                std: &NORMALIZE_STD ${oc.env:NORMALIZE_STD, [0.229, 0.224, 0.225]}
+        - class_path: eva.vision.datasets.MoNuSAC
+          init_args:
+            <<: *PREDICT_DATASET_ARGS
+            split: test
+    dataloaders:
+      train:
+        batch_size: &BATCH_SIZE ${oc.env:BATCH_SIZE, 64}
+        shuffle: true
+      val:
+        batch_size: *BATCH_SIZE
+      predict:
+        batch_size: &PREDICT_BATCH_SIZE ${oc.env:PREDICT_BATCH_SIZE, 64}

configs/vision/dino_vit/offline/total_segmentator_2d.yaml

@@ -0,0 +1,119 @@
+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/${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224}/total_segmentator_2d}
+    max_steps: &MAX_STEPS ${oc.env:MAX_STEPS, 20000}
+    callbacks:
+      - class_path: eva.callbacks.ConfigurationLogger
+      - class_path: eva.vision.callbacks.SemanticSegmentationLogger
+        init_args:
+          log_every_n_steps: 1000
+          log_images: false
+      - class_path: eva.callbacks.SegmentationEmbeddingsWriter
+        init_args:
+          output_dir: &DATASET_EMBEDDINGS_ROOT ${oc.env:EMBEDDINGS_ROOT, ./data/embeddings}/${oc.env:DINO_BACKBONE, dino_vits16}/total_segmentator_2d
+          dataloader_idx_map:
+            0: train
+            1: val
+            2: test
+          metadata_keys: ["slice_index"]
+          backbone:
+            class_path: eva.vision.models.networks.encoders.TimmEncoder
+            init_args:
+              model_name: ${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224_dino}
+              pretrained: ${oc.env:MODEL_PRETRAINED, true}
+              out_indices: ${oc.env:TIMM_MODEL_OUT_INDICES, 1}
+              model_arguments:
+                dynamic_img_size: true
+    logger:
+      - class_path: lightning.pytorch.loggers.TensorBoardLogger
+        init_args:
+          save_dir: *OUTPUT_ROOT
+          name: ""
+model:
+  class_path: eva.vision.models.modules.SemanticSegmentationModule
+  init_args:
+    decoder:
+      class_path: eva.vision.models.networks.decoders.segmentation.ConvDecoder
+      init_args:
+        layers:
+          class_path: torch.nn.Conv2d
+          init_args:
+            in_channels: ${oc.env:DECODER_IN_FEATURES, 384}
+            out_channels:  &NUM_CLASSES 118
+            kernel_size: [1, 1]
+    criterion: torch.nn.CrossEntropyLoss
+    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.EmbeddingsSegmentationDataset
+        init_args: &DATASET_ARGS
+          root: *DATASET_EMBEDDINGS_ROOT
+          manifest_file: manifest.csv
+          split: train
+      val:
+        class_path: eva.vision.datasets.EmbeddingsSegmentationDataset
+        init_args:
+          <<: *DATASET_ARGS
+          split: val
+      predict:
+        - class_path: eva.vision.datasets.TotalSegmentator2D
+          init_args: &PREDICT_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.ResizeAndCrop
+              init_args:
+                size: ${oc.env:RESIZE_DIM, 224}  
+                mean: &NORMALIZE_MEAN ${oc.env:NORMALIZE_MEAN, [0.485, 0.456, 0.406]} 
+                std: &NORMALIZE_STD ${oc.env:NORMALIZE_STD, [0.229, 0.224, 0.225]}
+        - class_path: eva.vision.datasets.TotalSegmentator2D
+          init_args:
+            <<: *PREDICT_DATASET_ARGS
+            split: val
+        - class_path: eva.vision.datasets.TotalSegmentator2D
+          init_args:
+            <<: *PREDICT_DATASET_ARGS
+            split: test
+    dataloaders:
+      train:
+        batch_size: &BATCH_SIZE ${oc.env:BATCH_SIZE, 64}
+        shuffle: true
+      val:
+        batch_size: *BATCH_SIZE
+      test:
+        batch_size: *BATCH_SIZE
+      predict:
+        batch_size: &PREDICT_BATCH_SIZE ${oc.env:PREDICT_BATCH_SIZE, 64}

configs/vision/dino_vit/online/total_segmentator_2d.yaml

@@ -10,8 +10,8 @@ trainer:
       - 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]
+          mean: &NORMALIZE_MEAN ${oc.env:NORMALIZE_MEAN, [0.485, 0.456, 0.406]} 
+          std: &NORMALIZE_STD ${oc.env:NORMALIZE_STD, [0.229, 0.224, 0.225]}
       - class_path: lightning.pytorch.callbacks.ModelCheckpoint
         init_args:
           filename: best
@@ -36,8 +36,8 @@ model:
     encoder:
       class_path: eva.vision.models.networks.encoders.TimmEncoder
       init_args:
-        model_name: ${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224}
-        pretrained: true
+        model_name: ${oc.env:TIMM_MODEL_NAME, vit_small_patch16_224_dino}
+        pretrained: ${oc.env:MODEL_PRETRAINED, true}
         out_indices: ${oc.env:TIMM_MODEL_OUT_INDICES, 1}
         model_arguments:
           dynamic_img_size: true