This repository contains all code related to the experimental part of my MSc. thesis on Semantic Segmentation of Medical Images using Deep Learning. The thesis focuses on the topic of processing 3D medical images with 2D Convolutional Neural Networks alternatively to their 3D counterparts.

I focused on ACS convolutions and assessed the impact of neural architecural elements that match the way this operator works. With ACS convolutions, any 2D Convolutional Neural Network can directly process 3D images by splitting the filters in each convolutional layer into 3 independent sets. Each set of filters is responsible for modelling a given axis of the input given to the convolutional layer. The original implementation for ACS convolutions can be found in this repository.

• An ACS U-Net is either equivalent or superior to an equivalent 3D U-Net in terms of segmentation performance.
• Group Normalization (with three channel groups) enables superior segmentation performance for ACS models.
• Group Normalization makes the value distributions of axis-specific feature maps more distinct amogst them when compared to Batch Normalization.
• Through multi-task learning, adding an auxiliary classfication task to distinguish between feature maps of specific axes does not have an impact on segmentation performance.

All commands are executed through a command line interface present in pytorch/main.py. It consists of a generic tool to run experiments and benchmark CNN models for semantic segmentation of 3D images.

First, create a python environment (Python 3.7+) and install the dependencies with:

pip install -r requirements.txt

Due to the nature of 3D images, it is highly recommended (if not essential) to have a GPU available, preferably with at least 16GB of RAM.

To run experiments you must provide the following arguments to the CLI:

• -c: command name. To perform a model training run with randomly initialized weights, the command name is from_scratch_supervised. This is the main command for training models. Other relevant commands will be described later on.
• -d: dataset. This must correspond to a key of the dataset_map dictionary of pytorch/utils.py. This dictionary defines the name of the dataset and the location of its files.
• --model: For all available models, check the _load_model method of the Trainer class in pytorch/finetune.py.

Additionally, there are optional arguments that may need to be specified:

• --num_cv_folds: Controls the use of cross validation in the experiment. If not set, cross validation won´t be used. You must specify an integer which corresponds to the amount of CV-folds to be considered.
• --in_channels and --out_channels: both default to 1. These correspond to the input and output channels of the first and last convolutional layer of the CNN models available. You must change them according to the dataset used. Set --in_channels to the number of channels/modalities present in the images and --out_channels to the number of target segmentation classes + 1 (as the background is also considered a target class).

You can also change certain aspects/hyperparameters of training such as the learning rate, loss function, batch size, early stop patience, scheduler and more by specifying the corresponding argument in the CLI. Most of the arguments end with a _sup suffix to indicate that they correspond to the task of supervised semantic segmentation. In the beggining of the project, I experimented with self-supervised (ss) pretraining of models so this distinction was made.

The hyperparameters and corresponding argument names as used in the CLI are:

An example of a training run would thus be:

python pytorch/main.py -c from_scratch_supervised -d LIDC --model unet_3d --in_channels 1 --out_channels 2 --batch_size_sup 22 --patience_sup_terminate 25 --patience_sup 10 --num_cv_folds 5

Keep in mind that this corresponds to a single training run. For example, if you are using 5-fold cross-validation, this command would have to be run 5 times, one for each individual run.

Depending on the parameters you have provided, a task_name will be assigned to the experiment you are conducting. All relevant files from a given a training run are stored automatically in directories under their task_name.

For example, with the provided command example, the task name would be: FROM_SCRATCH_LIDC_UNET_3D/ and each run y (out of 5) would have its own subfolder in the form of FROM_SCRATCH_LIDC_UNET_3D/run_y/

The task name thus contains relevant aspects of training such that you can easily identify your experiment. You can visualize the task_name generated during training.

At the end of each run, you can visualize convergence curves and other aspects of training by launching tensorboard. Tensorboard data is stored in the runs/ folder. To visualize the tensorboard logs of a given task_namerun:

tensorboard --logdir=runs/your_task_name

Once models are done training, obtaining segmentation performance is a 2 stage process:

• Run python pytorch/main.py -c test. This will generate .json files in the test_results/ folder containing dice coefficient and hausdorff distance metrics on the test set for each model run. If you're in a hurry or just interested in a particular experiment you can additionally pass the --task_name argument to match the task_name you want to compute the results on.
• Analyse/visualize results with the results.ipynb jupyter notebook. The results for all runs are parsed and displayed in tables. You get access to things like average metrics and standard errors. You can also perform statistical tests between different task_names to assess statistical significance in difference of results.

For more information on how test metrics are computed check chapter 4 of the MSc. thesis.

If you want to get image examples of segmentations, run the command:

python pytorch/main.py -c save_images

Once more, you can specify task_name to only get segmentation examples for a given run. Images are stored in the viz_samples/ directory. Ground truth segmentations are displayed along the models predicted segmentation.

• Make sure the dataset_map dictionary in pytorch/utils.py has a key corresponding to your dataset name (as you would pass with the -d dataset flag when performing a training run). Its value should be the directory where the dataset is located.
• The directory is supposed to have a structure with two folders: x/and y/. As in:

dataset_directory
└─── x/
└─── y/    

• x/ should contain images and y/ their corresponding segmentation masks. The name of an image and its corresponding segmentation should be the exact same.

Only .npy files are accepted. Also, each file is supposed to correspond to 1 image. The .npy file (numpy array) of a 3D image should have 4 dimensions: $$ Channels x H x D x W $$

• Create your network architecture.
• Import it in pytorch/finetune.py and add it to the _load_model method.
• Add the model name as an allowed one in pytorch/utils.py.

In classic thesis style, there are many unused pieces of code resulting from failures or ideas that did not get to be implemented or used much. The main functionality of the CLI has been described. Feel free to open an issue or contact me if you have further questions or something is not working as expected.