Configurable augmentations (#61)

* configurable augmentations

* increase example noise std

* update example configs

* update dlmbl script

* update scripts
removed auto-generated notebooks

* fix import path to use the redefined transforms

examples/configs/fit_example.yml

@@ -2,42 +2,21 @@
 seed_everything: true
 trainer:
   accelerator: auto
-  strategy: auto # ddp_find_unused_parameters_true for more GPUs
+  strategy: auto
   devices: auto
   num_nodes: 1
   precision: 32-true
   callbacks: null
   fast_dev_run: false
   max_epochs: null
-  min_epochs: null
-  max_steps: -1
-  min_steps: null
-  max_time: null
-  limit_train_batches: null
-  limit_val_batches: null
-  limit_test_batches: null
-  limit_predict_batches: null
-  overfit_batches: 0.0
   val_check_interval: null
   check_val_every_n_epoch: 1
-  num_sanity_val_steps: null
   log_every_n_steps: null
   enable_checkpointing: null
-  enable_progress_bar: null
-  enable_model_summary: null
-  accumulate_grad_batches: 1
-  gradient_clip_val: null
-  gradient_clip_algorithm: null
-  deterministic: null
-  benchmark: null
   inference_mode: true
   use_distributed_sampler: true
-  profiler: null
-  detect_anomaly: false
-  barebones: false
   plugins: null
   sync_batchnorm: true
-  reload_dataloaders_every_n_epochs: 0
   default_root_dir: null
 model:
   model_config: {}
@@ -47,17 +26,52 @@ model:
   log_batches_per_epoch: 8
   log_samples_per_batch: 1
 data:
-  data_path: null
-  source_channel: null
-  target_channel: null
-  z_window_size: null
-  split_ratio: null
-  batch_size: 16
-  num_workers: 8
-  yx_patch_size:
-  - 256
-  - 256
-  augment: true
+  data_path: /path/to/data.hcs.ome.zarr
+  source_channel: source
+  target_channel:
+    - target_1
+    - target_2
+  z_window_size: 5
+  split_ratio: 0.8
+  batch_size: 32
+  num_workers: 16
+  yx_patch_size: [256, 256]
+  augmentations:
+    - class_path: viscy.transforms.RandWeightedCropd
+      init_args:
+        keys: [source, target_1, target_2]
+        w_key: target_1
+        spatial_size: [-1, 512, 512]
+        num_samples: 1
+    - class_path: viscy.transforms.RandAffined
+      init_args:
+        keys: [source, target_1, target_2]
+        prob: 0.5
+        rotate_range: [3.14, 0.0, 0.0]
+        shear_range: [0.0, 0.05, 0.05]
+        scale_range: [0.2, 0.3, 0.3]
+    - class_path: viscy.transforms.RandAdjustContrastd
+      init_args:
+        keys: source
+        prob: 0.3
+        gamma: [0.75, 1.5]
+    - class_path: viscy.transforms.RandScaleIntensityd
+      init_args:
+        keys: source
+        prob: 0.5
+        factors: 0.5
+    - class_path: viscy.transforms.RandGaussianNoised
+      init_args:
+        keys: source
+        prob: 0.5
+        mean: 0.0
+        std: 1.0
+    - class_path: viscy.transforms.RandGaussianSmoothd
+      init_args:
+        keys: source
+        prob: 0.5
+        sigma_z: [0.25, 1.5]
+        sigma_y: [0.25, 1.5]
+        sigma_x: [0.25, 1.5]
   caching: false
-  normalize_source: false
-ckpt_path: null
+  normalize_source: true

examples/configs/predict_example.yml

@@ -61,7 +61,6 @@ predict:
     yx_patch_size:
       - 256
       - 256
-    augment: true
     caching: false
     normalize_source: false
     predict_scale_source: null

examples/configs/test_example.yml

@@ -60,14 +60,8 @@ data:
   yx_patch_size:
   - 256
   - 256
-  augment: true
   caching: false
   normalize_source: false
   ground_truth_masks: null
-  train_z_scale_range:
-  - 0.0
-  - 0.0
-  train_noise_std: 0.0
-  train_patches_per_stack: 1
 ckpt_path: null
 verbose: true

examples/demo_dlmbl/debug_log_graph.py

@@ -80,7 +80,7 @@
     num_workers=8,
     architecture="2D",
     yx_patch_size=YX_PATCH_SIZE,
-    augment=True,
+    augmentations=None,
 )
 phase2fluor_data.setup("fit")