I have moved from unbatched transformations to adding the possibility of applying the pipeline to B segmentations, obtaining B different images, each with separate sampled parameters. Of course, this requires that all input segmentations have the same shape and orientation (due to FlipTransform). The only change in the usage of the code is that the outputted images have an extra dimension, i.e., instead of being [C H W D] they will be [B C H W D], where C is a user-specified number of channels that you want to generate for a given anatomy.
This repository implements a pytorch-based version of SynthSeg/Labels to Image as described in Billot et al. Note that such an implementation has already been done – and probably better – by torchio (this is research code and not guaranteed to work).
However, this repository is motivated by the following points:
- To side-step
tensorflow(which I am personally not a fan of). - Adopting
batchgeneratorsv2's more lightweight approach to transformations. - To have more control over the pipeline (understand it better), and also take the opportunity to learn.
This code can be run on the GPU or CPU either when loading batches or on already batched segmentations (as long as they have the same spatial shapes and orientations) for training a neural net for medical image segmentation. Essentially, the main "entrypoints" are labels_to_image.py and _gen_params.py. The former creates the pipeline using the parameters defined in the latter.
To install this package use:
git clone https://github.com/vcasellesb/lab2im-torch.git
cd lab2im-torch
pip install (-e) .Please note that to load segmentations another package will be required, such as nibabel or SimpleITK (as evidenced in the following section and in the sample script (test_script.py)), which are not included in this package's requirements.
Please refer to the sample test_script.py for a little sample. Running it will require installation of nibabel (pip install nibabel). In short, to run the pipeline on a segmentation, you should pass it to its __call__ method as a keyword argument, as follows:
from lab2im import LabelsToImage, GenerationParams
label = nib.load(...)
# Important: pass segmentations as 4D pytorch tensors
data = torch.from_numpy(np.asarray(label.dataobj))[None]
# Also, either reorient your image to RAS, or pass the following parameter to the pipeline.
# It is used to determine which axis to do flips on (since the human body tends to be only
# symmetric in the R/L axis)
axis_codes = nib.aff2axcodes(label.affine)
input_resolution = label.header.get_zooms()
lab2im = LabelsToImage(GenerationParams())
data = lab2im(**{'seg': data, 'input_resolution': input_resolution, 'axis_codes': axis_codes})That is, each transform gets the data (case, subject in torchio's parlance) as a dictionary, and stores the image and modifies the segmentation using the following keys:
image_key = 'im'
segmentation_key = 'seg'Although this is subject to change.