Tune augmentations with CLI and config for contrastive models (#126)

* fix initial crop size

* make representation a sub-package

* re-export ScaleIntensityRangePercentilesd
for jsonargparse

* add cli module for contrastive models with triplet data

* update docstring

Author: ziw-liu

viscy/cli/contrastive_triplet.py

@@ -0,0 +1,41 @@
+import logging
+import os
+from datetime import datetime
+
+import torch
+from jsonargparse import lazy_instance
+from lightning.pytorch.cli import LightningCLI
+from lightning.pytorch.loggers import TensorBoardLogger
+
+from viscy.data.triplet import TripletDataModule
+from viscy.light.engine import ContrastiveModule
+
+
+class ContrastiveLightningCLI(LightningCLI):
+    """Lightning CLI with default logger."""
+
+    def add_arguments_to_parser(self, parser):
+        parser.set_defaults(
+            {
+                "trainer.logger": lazy_instance(
+                    TensorBoardLogger,
+                    save_dir="",
+                    version=datetime.now().strftime(r"%Y%m%d-%H%M%S"),
+                    log_graph=True,
+                )
+            }
+        )
+
+
+def main():
+    """Main Lightning CLI entry point."""
+    log_level = os.getenv("VISCY_LOG_LEVEL", logging.INFO)
+    logging.getLogger("lightning.pytorch").setLevel(log_level)
+    torch.set_float32_matmul_precision("high")
+    _ = ContrastiveLightningCLI(
+        model_class=ContrastiveModule, datamodule_class=TripletDataModule
+    )
+
+
+if __name__ == "__main__":
+    main()

viscy/data/triplet.py

@@ -262,7 +262,7 @@ def _setup_fit(self, dataset_settings: dict):
         self.train_dataset = TripletDataset(
             positions=train_positions,
             tracks_tables=train_tracks_tables,
-            initial_yx_patch_size=self.yx_patch_size,
+            initial_yx_patch_size=self.initial_yx_patch_size,
             anchor_transform=no_aug_transform,
             positive_transform=augment_transform,
             negative_transform=augment_transform,
@@ -273,7 +273,7 @@ def _setup_fit(self, dataset_settings: dict):
         self.val_dataset = TripletDataset(
             positions=val_positions,
             tracks_tables=val_tracks_tables,
-            initial_yx_patch_size=self.yx_patch_size,
+            initial_yx_patch_size=self.initial_yx_patch_size,
             anchor_transform=no_aug_transform,
             positive_transform=augment_transform,
             negative_transform=augment_transform,

viscy/representation/__init__.py

@@ -11,6 +11,7 @@
     RandomizableTransform,
     RandScaleIntensityd,
     RandWeightedCropd,
+    ScaleIntensityRangePercentilesd,
 )
 from monai.transforms.transform import Randomizable
 from numpy.random.mtrand import RandomState as RandomState
@@ -127,6 +128,34 @@ def __init__(
         )
 
 
+class ScaleIntensityRangePercentilesd(ScaleIntensityRangePercentilesd):
+    def __init__(
+        self,
+        keys: Union[Sequence[str], str],
+        lower: float,
+        upper: float,
+        b_min: float | None,
+        b_max: float | None,
+        clip: bool = False,
+        relative: bool = False,
+        channel_wise: bool = False,
+        dtype: Union[Sequence[str], str] = None,
+        allow_missing_keys: bool = False,
+    ):
+        super().__init__(
+            keys=keys,
+            lower=lower,
+            upper=upper,
+            b_min=b_min,
+            b_max=b_max,
+            clip=clip,
+            relative=relative,
+            channel_wise=channel_wise,
+            dtype=dtype,
+            allow_missing_keys=allow_missing_keys,
+        )
+
+
 class NormalizeSampled(MapTransform):
     """
     Normalize the sample