Pixel-wise segmentation for biomedical images using Pytorch. This project was part of my Bachelor Thesis.

The code was tested on a biomedical breast cancer dataset which can't be provided due to data privacy. A very similar dataset can be found at bioimage.ucsb.edu Note: The images have a different size than the images I'm using. The network achitecture (reshaping of tensors, padding after each convolution) needs to be tuned a little bit to work with the data.

I have implemented the following architecture:

• U-Net

torch
torchvision
numpy
PIL
pydensecrf

The code is based on Python 3.6

.
├── data                    # training dataset
├── val                     # validation set
├── checkpoints             # checkpoints to store the model

Data and Val folder needs to contain .txt files with filenames of the data. The "create_txt" script can be used to generate these.

There are multiple parameters which can be used:

main.py

-e	# number of epochs to train
-l	# learning rate
-g	# use this parameter to utilize your GPU
-c	# load a pretrained model

predict.py

-m	# path to the pretrained model / checkpoint
-i	# input image to predict
-o	# filename of the output image
-c	# GPU support is enabled by default. Use this parameter to predict on CPU

There are some more parameters which can be useful.

Further evaluation needs to be done in the future. For now the net reached an Accuracy of 87% (Dice Coefficient of 0.87) on a small validation set.

I'm using a Nvidia GTX 1070 (8GB VRAM) to train the net. The highest memory usage I could observe was about 5,5GB. The size of the images is 510x512 which will be downscaled to 255x256. Downscaling even further will lower the memory needed.

Training took about 1 1/2 hours.