Skip to content

Commit

Permalink
Feature/head cleanup (#336)
Browse files Browse the repository at this point in the history
* only run heads that dont have existing outputs

* make gan subclass of multitask model

* update saving code

* remove conv_proj_layer

* make superres head class gan and resblockshead

* remove metric calculation

* update configs

* update readme

* remove conv_proj_layer;

* precommit

---------

Co-authored-by: Benjamin Morris <[email protected]>
  • Loading branch information
benjijamorris and Benjamin Morris authored Feb 23, 2024
1 parent 2b89383 commit 4860c8b
Show file tree
Hide file tree
Showing 24 changed files with 772 additions and 310 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -29,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 using conditional GANs (e.g. predicting 100x images from 20x images), semantic and instance segmentation, and label-free prediction. We also provide configs for Masked Autoencoder (MAE) pretraining using a Vision Transformer (ViT) backbone (currently only supported for 3D images) and for training segmentation decoders from these pretrained features. Representation learning code includes a wide variety of Variational Auto Encoder (VAE) architectures. Due to dependency issues, equivariant autoencoders are not currently supported on Windows.
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 using conditional GANs (e.g. predicting 100x images from 20x images), semantic and instance segmentation, and label-free prediction. We also provide configs for Masked Autoencoder (MAE) pretraining using a Vision Transformer (ViT) backbone and for training segmentation decoders from these pretrained features. Representation learning code includes a wide variety of Variational Auto Encoder (VAE) architectures. Due to dependency issues, equivariant autoencoders are not currently supported on Windows.

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.

Expand Down
205 changes: 205 additions & 0 deletions configs/data/im2im/gan_superres.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,205 @@
_target_: cyto_dl.datamodules.dataframe.DataframeDatamodule

path:
cache_dir:

num_workers: 0
batch_size: 1
pin_memory: True
split_column:
columns:
- ${source_col}
- ${target_col}

transforms:
train:
_target_: monai.transforms.Compose
transforms:
# 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
# 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: ${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: ${source_col}
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]
- _target_: monai.transforms.ScaleIntensityRangePercentilesd
keys: ${target_col}
lower: 0.01
upper: 99.99
b_min: -1
b_max: 1
clip: True
# input to synthetic image generation model is a semantic segmentation
- _target_: monai.transforms.ThresholdIntensityd
keys: ${source_col}
threshold: 0.1
above: False
cval: 1
- _target_: cyto_dl.image.transforms.RandomMultiScaleCropd
keys: ${data.columns}
patch_shape: ${data._aux.patch_shape}
patch_per_image: 1
scales_dict: ${kv_to_dict:${data._aux._scales_dict}}

test:
_target_: monai.transforms.Compose
transforms:
# 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: ${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: ${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}
#crop image so that it is divisible by 4 for downsampling
- _target_: monai.transforms.CenterSpatialCropd
keys: ${data.columns}
roi_size: [64, -1, -1]
- _target_: monai.transforms.Zoomd
keys: ${source_col}
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]
- _target_: monai.transforms.ScaleIntensityRangePercentilesd
keys: ${target_col}
lower: 0.01
upper: 99.99
b_min: -1
b_max: 1
clip: True
# input to synthetic image generation model is a semantic segmentation
- _target_: monai.transforms.ThresholdIntensityd
keys: ${source_col}
threshold: 0.1
above: False
cval: 1

predict:
_target_: monai.transforms.Compose
transforms:
# 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: ${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}
#crop image so that it is divisible by 4 for downsampling
- _target_: monai.transforms.CenterSpatialCropd
keys: ${source_col}
roi_size: [64, -1, -1]
- _target_: monai.transforms.Zoomd
keys: ${source_col}
zoom: 0.25
keep_size: False

- _target_: monai.transforms.ToTensord
keys: ${source_col}
# input to synthetic image generation model is a semantic segmentation
- _target_: monai.transforms.ThresholdIntensityd
keys: ${source_col}
threshold: 0.1
above: False
cval: 1

valid:
_target_: monai.transforms.Compose
transforms:
# 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: ${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: ${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: ${source_col}
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]
- _target_: monai.transforms.ScaleIntensityRangePercentilesd
keys: ${target_col}
lower: 0.01
upper: 99.99
b_min: -1
b_max: 1
clip: True
# input to synthetic image generation model is a semantic segmentation
- _target_: monai.transforms.ThresholdIntensityd
keys: ${source_col}
threshold: 0.1
above: False
cval: 1
- _target_: cyto_dl.image.transforms.RandomMultiScaleCropd
keys: ${data.columns}
patch_shape: ${data._aux.patch_shape}
patch_per_image: 1
scales_dict: ${kv_to_dict:${data._aux._scales_dict}}

_aux:
_scales_dict:
- - ${target_col}
- [4]
- - ${source_col}
- [1]
Loading

0 comments on commit 4860c8b

Please sign in to comment.