Bugfix/2d tests (AllenCellModeling#307)

* allow 3d/2d without config changes * standardize, add baseic 2d/3d patch shapes * update postprocessing and inference params * remove dimensions where value is None, needed for 2d/3d data configs * make compile check stricter * fix dtypes for iseg * make gan work in 2d * test 2d and 3d * clarify what default configs are for * increase timeout --------- Co-authored-by: Benjamin Morris <[email protected]>
colobas · Nov 3, 2023 · 53d2c5d · 53d2c5d
1 parent 4c5c8ca
commit 53d2c5d
Show file tree

Hide file tree

Showing 24 changed files with 252 additions and 110 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -13,7 +13,7 @@ concurrency:
 jobs:
   test:
     runs-on: ${{ matrix.os }}
-    timeout-minutes: 20
+    timeout-minutes: 70
     strategy:
       fail-fast: false
       matrix:

diff --git a/README.md b/README.md
@@ -5,6 +5,7 @@
   <source media="(prefers-color-scheme: light)" srcset="https://github.com/AllenCellModeling/cyto-dl/blob/b73e6f357727e3b42adea8540c86f2475ea60379/docs/CytoDL-logo-1C-onLight.png">
   <img src="https://github.com/AllenCellModeling/cyto-dl/blob/b73e6f357727e3b42adea8540c86f2475ea60379/docs/CytoDL-logo-1C-onLight.png">
 </picture> -->
+
 <h1>CytoDL</h1>
 
 <a href="https://pytorch.org/get-started/locally/"><img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-ee4c2c?logo=pytorch&logoColor=white"></a>
@@ -28,7 +29,7 @@ As part of the [Allen Institute for Cell Science's](allencell.org) mission to un
 
 The bulk of `CytoDL`'s underlying structure bases the [lightning-hydra-template](https://github.com/ashleve/lightning-hydra-template) organization - we highly recommend that you familiarize yourself with their (short) docs for detailed instructions on running training, overrides, etc.
 
-Our currently available code is roughly split into two domains: image-to-image transformations and representation learning. The image-to-image code (denoted im2im) contains configuration files detailing how to train and predict using models for resolution enhancement (e.g. predicting 100x images from 20x images), semantic and instance segmentation, and label-free prediction. Representation learning code includes a wide variety of Variational Auto Encoder (VAE) architectures.
+Our currently available code is roughly split into two domains: image-to-image transformations and representation learning. The image-to-image code (denoted im2im) contains configuration files detailing how to train and predict using models for resolution enhancement (e.g. predicting 100x images from 20x images), semantic and instance segmentation, and label-free prediction. Representation learning code includes a wide variety of Variational Auto Encoder (VAE) architectures. Note that these default models are very small and by default run on heavily downsampled data in order to make tests run efficiently - for best performance, the model size should be increased and downsampling removed.
 
 As we rely on recent versions of pytorch, users wishing to train and run models on GPU hardware will need up-to-date NVIDIA drivers. Users with older GPUs should not expect code to work out of the box. Similarly, we do not currently support training/predicting on Mac GPUs. In most cases, cpu-based training should work when GPU training fails.
 

diff --git a/configs/data/im2im/gan.yaml b/configs/data/im2im/gan.yaml
@@ -15,28 +15,33 @@ transforms:
   train:
     _target_: monai.transforms.Compose
     transforms:
-      # channels are [blank, membrane,blank, structure, nuclear dye, brightfield ]
+      # channels are [blank, membrane,blank, structure, blank, nuclear dye, brightfield ]
       # target is the nuclear dyeimage
       - _target_: monai.transforms.LoadImaged
         keys: ${target_col}
         reader:
-          - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
-            C: 4
+          - _target_:
+              cyto_dl.image.io.MonaiBioReader
+              # NOTE: eval is used so only the experiment file is required to change for beginning users. This is not recommended when creating your own configs.
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
+            C: 5
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
       # channels are [nucseg, cellseg, nuclear boundary seg, cell boundary seg]
       # source image is the segmentation
       - _target_: monai.transforms.LoadImaged
         keys: ${source_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
             C: 0
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
       - _target_: monai.transforms.EnsureChannelFirstd
         channel_dim: "no_channel"
         keys: ${data.columns}
       - _target_: monai.transforms.Zoomd
         keys: ${data.columns}
         zoom: 0.25
+        keep_size: False
       - _target_: monai.transforms.ToTensord
         keys: ${data.columns}
       # GANs use Tanh as final activation, target has to be in range [-1,1]
@@ -62,28 +67,31 @@ transforms:
   test:
     _target_: monai.transforms.Compose
     transforms:
-      # channels are [blank, membrane,blank, structure, nuclear dye, brightfield ]
+      # channels are [blank, membrane,blank, structure, blank, nuclear dye, brightfield ]
       # target is the nuclear dyeimage
       - _target_: monai.transforms.LoadImaged
         keys: ${target_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
-            C: 4
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
+            C: 5
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
       # channels are [nucseg, cellseg, nuclear boundary seg, cell boundary seg]
       # source image is the segmentation
       - _target_: monai.transforms.LoadImaged
         keys: ${source_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
             C: 0
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
       - _target_: monai.transforms.EnsureChannelFirstd
         channel_dim: "no_channel"
         keys: ${data.columns}
       - _target_: monai.transforms.Zoomd
         keys: ${data.columns}
         zoom: 0.25
+        keep_size: False
       - _target_: monai.transforms.ToTensord
         keys: ${data.columns}
       # GANs use Tanh as final activation, target has to be in range [-1,1]
@@ -110,14 +118,17 @@ transforms:
         keys: ${source_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
             C: 0
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
       - _target_: monai.transforms.EnsureChannelFirstd
         channel_dim: "no_channel"
         keys: ${source_col}
       - _target_: monai.transforms.Zoomd
         keys: ${data.columns}
         zoom: 0.25
+        keep_size: False
+
       - _target_: monai.transforms.ToTensord
         keys: ${source_col}
       # input to synthetic image generation model is a semantic segmentation
@@ -130,28 +141,32 @@ transforms:
   valid:
     _target_: monai.transforms.Compose
     transforms:
-      # channels are [blank, membrane,blank, structure, nuclear dye, brightfield ]
+      # channels are [blank, membrane,blank, structure, blank, nuclear dye, brightfield ]
       # target is the nuclear dyeimage
       - _target_: monai.transforms.LoadImaged
         keys: ${target_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
-            C: 4
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
+            C: 5
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
       # channels are [nucseg, cellseg, nuclear boundary seg, cell boundary seg]
       # source image is the segmentation
       - _target_: monai.transforms.LoadImaged
         keys: ${source_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
             C: 0
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
       - _target_: monai.transforms.EnsureChannelFirstd
         channel_dim: "no_channel"
         keys: ${data.columns}
       - _target_: monai.transforms.Zoomd
         keys: ${data.columns}
         zoom: 0.25
+        keep_size: False
+
       - _target_: monai.transforms.ToTensord
         keys: ${data.columns}
       # GANs use Tanh as final activation, target has to be in range [-1,1]
@@ -175,7 +190,6 @@ transforms:
         scales_dict: ${kv_to_dict:${data._aux._scales_dict}}
 
 _aux:
-  patch_shape: [32, 512, 512]
   _scales_dict:
     - - ${target_col}
       - [1]

diff --git a/configs/data/im2im/instance_seg.yaml b/configs/data/im2im/instance_seg.yaml
@@ -3,7 +3,7 @@ _target_: cyto_dl.datamodules.dataframe.DataframeDatamodule
 path:
 cache_dir:
 
-num_workers: 0
+num_workers: 1
 batch_size: 1
 pin_memory: True
 split_column:
@@ -15,49 +15,60 @@ transforms:
   train:
     _target_: monai.transforms.Compose
     transforms:
-      # channels are [blank, membrane,blank, structure, nuclear dye, brightfield ]
+      # channels are [blank, membrane,blank, structure, blank, nuclear dye, brightfield ]
       - _target_: monai.transforms.LoadImaged
         keys: ${source_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
-            C: 4
+            # NOTE: eval is used so only the experiment file is required to change for beginning users. This is not recommended when creating your own configs.
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
+            C: 5
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
       # channels are [nucseg, cellseg, nuclear boundary seg, cell boundary seg]
       - _target_: monai.transforms.LoadImaged
         keys: ${target_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
             C: 0
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
+      # transforms require CZYX images, but we are loading as ZYX
       - _target_: monai.transforms.EnsureChannelFirstd
-        channel_dim: "no_channel"
         keys: ${data.columns}
+        channel_dim: "no_channel"
+      # 4x downsample to speed up default
       - _target_: monai.transforms.Zoomd
-        keys: ${data.columns}
+        keys:
+          - ${source_col}
+          - ${target_col}
         zoom: 0.25
+        keep_size: False
+        mode: ["area", "nearest"]
+      # preprocess instance segmentation into vector field
       - _target_: cyto_dl.models.im2im.utils.InstanceSegPreprocessd
         label_keys: ${target_col}
         dim: ${spatial_dims}
       - _target_: monai.transforms.ToTensord
         keys: ${data.columns}
+      # z-score normalization
       - _target_: monai.transforms.NormalizeIntensityd
         keys: ${source_col}
         channel_wise: True
+      # randomly select patches for training
       - _target_: cyto_dl.image.transforms.RandomMultiScaleCropd
         keys: ${data.columns}
         patch_shape: ${data._aux.patch_shape}
-        patch_per_image: 1
+        patch_per_image: 4
         scales_dict: ${kv_to_dict:${data._aux._scales_dict}}
+      # augmentation transforms
       - _target_: monai.transforms.RandHistogramShiftd
         prob: 0.1
         keys: ${source_col}
         num_control_points: [90, 500]
-
       - _target_: monai.transforms.RandStdShiftIntensityd
         prob: 0.1
         keys: ${source_col}
         factors: 0.1
-
       - _target_: monai.transforms.RandAdjustContrastd
         prob: 0.1
         keys: ${source_col}
@@ -66,26 +77,32 @@ transforms:
   test:
     _target_: monai.transforms.Compose
     transforms:
-      # channels are [blank, membrane,blank, structure, nuclear dye, brightfield ]
+      # channels are [blank, membrane,blank, structure, blank, nuclear dye, brightfield ]
       - _target_: monai.transforms.LoadImaged
         keys: ${source_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
             C: 5
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
       # channels are [nucseg, cellseg, nuclear boundary seg, cell boundary seg]
       - _target_: monai.transforms.LoadImaged
         keys: ${target_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
             C: 0
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
       - _target_: monai.transforms.EnsureChannelFirstd
         channel_dim: "no_channel"
         keys: ${data.columns}
       - _target_: monai.transforms.Zoomd
-        keys: ${data.columns}
+        keys:
+          - ${source_col}
+          - ${target_col}
         zoom: 0.25
+        keep_size: False
+        mode: ["area", "nearest"]
       - _target_: cyto_dl.models.im2im.utils.InstanceSegPreprocessd
         label_keys: ${target_col}
         dim: ${spatial_dims}
@@ -98,19 +115,24 @@ transforms:
   predict:
     _target_: monai.transforms.Compose
     transforms:
-      # channels are [blank, membrane,blank, structure, nuclear dye, brightfield ]
+      # channels are [blank, membrane,blank, structure, blank, nuclear dye, brightfield ]
       - _target_: monai.transforms.LoadImaged
         keys: ${source_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
             C: 5
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
       - _target_: monai.transforms.EnsureChannelFirstd
         channel_dim: "no_channel"
         keys: ${data.columns}
       - _target_: monai.transforms.Zoomd
-        keys: ${data.columns}
+        keys:
+          - ${source_col}
+          - ${target_col}
         zoom: 0.25
+        keep_size: False
+        mode: ["area", "nearest"]
       - _target_: monai.transforms.ToTensord
         keys: ${source_col}
       - _target_: monai.transforms.NormalizeIntensityd
@@ -120,27 +142,33 @@ transforms:
   valid:
     _target_: monai.transforms.Compose
     transforms:
-      # channels are [blank, membrane,blank, structure, nuclear dye, brightfield ]
+      # channels are [blank, membrane,blank, structure, blank, nuclear dye, brightfield ]
       - _target_: monai.transforms.LoadImaged
         keys: ${source_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
             C: 5
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
         expand_user: False
       # channels are [nucseg, cellseg, nuclear boundary seg, cell boundary seg]
       - _target_: monai.transforms.LoadImaged
         keys: ${target_col}
         reader:
           - _target_: cyto_dl.image.io.MonaiBioReader
-            dimension_order_out: "ZYX"
+            dimension_order_out: ${eval:'"ZYX" if ${spatial_dims}==3 else "YX"'}
             C: 0
+            Z: ${eval:'None if ${spatial_dims}==3 else 38'}
       - _target_: monai.transforms.EnsureChannelFirstd
         channel_dim: "no_channel"
         keys: ${data.columns}
       - _target_: monai.transforms.Zoomd
-        keys: ${data.columns}
+        keys:
+          - ${source_col}
+          - ${target_col}
         zoom: 0.25
+        keep_size: False
+        mode: ["area", "nearest"]
       - _target_: cyto_dl.models.im2im.utils.InstanceSegPreprocessd
         label_keys: ${target_col}
         dim: ${spatial_dims}