Overview

This program implements U-Net architectures [1-3] for binary segmentation, where a model is trained to map each pixel in an image to the range $[0,1]$. This was designed with transmission electron-microscopy images in mind, where the output pixel values are probabilities of belonging to a defect structure ($1$ is defect, $0$ is a non-defect).

This project implements four types of configurable architectures:

• U-Net
• U-Net++
• U-Net w/ EfficientNetB7 backbone
• U-Net++ w/ EfficientNetB7 backbone

Models can either be generated by training, or they can applied for inference. There is also a $k$-fold cross validation feature to study the bias associated with various validation sets.

The program operations and model architecture are configured via input files discussed in configs/README.md.

Use with Linux/WSL

If you intend to use Linux (such as on Penn State's Roar Collab cluster) or WSL with your own NVIDIA GPU, follow the installation instructions discussed next. If using an AMD Radeon GPU, ROCm will have to be used in place of CUDA. Note, a CPU can be used, however, training/inference takes an immensely long time.

Support for multiple GPUs has not been implemented. If you want this functionality, submit an issue and it will be discussed accordingly.

Use with Google Colab

If you cannot use Linux or do not have a powerful GPU with sufficient memory, Google Colab is a good alternative. To use it, open the unet_compare.ipynb notebook and follow the instructions in the first cell.

You can use a lightweight GPU for free; more powerful GPUs are only available through the purchase of computational units. Priority access is enabled through a subscription service.

Installation

Conda is used to contain the project and install the dependencies. If you don't have Conda already installed, it can be installed using Miniconda:

cd ~
curl -O https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh
# follow installation procedures
# restart shell to finish install

With Conda installed, the repository can be cloned and a virtual environment created.

git clone https://github.com/psu-rdmap/unet-compare.git
cd unet-compare
export CONDA_PKGS_DIRS=$PWD/env/conda_pkgs  
conda create -p env python=3.11
conda activate env/

Next, ensure install of Python and Pip (the package installer for Python) corresponds to the virtual environment.

which python   # /path/to/unet-compare/env/bin/python
which pip      # /path/to/unet-compare/env/bin/pip

Dependencies can then be installed.

conda install cudatoolkit==11.8.0
pip install --upgrade pip
pip install -r requirements.txt

The project should now be installed. The following line checks which devices Tensorflow recognizes:

python scripts/check_devices.py

Training & Inference

This section gives information about running the source code. The first step is to create a dataset with the correct structure.

1. Creating a dataset

Two training datasets are supplied in the data/ directory. Each dataset has the following structure:

data/dataset_name
├── images
│   ├── 1.ext
│   ├── 2.ext
│   └── ...
└── annotations
    ├── 1.ext
    ├── 2.ext
    └── ...

This format is for training and cross-validation. For inference, images can be directly placed in the dataset directory:

data/dataset_name
├── 1.ext
├── 2.ext
└── ...

Images are associated with their corresponding annotations by giving them the same filenames. All images or annotations must use a consistent file format. For example all images could be .jpg files, while all annotations could be .png files.

If you do not want to copy your own dataset into data/, it can be symlinked (i.e. create a shortcut) instead via

ln -s /path/to/dataset/ /path/to/unet-compare/data/

2. Running the program

The next step is to create a configs input file in the configs/ following the instructions in the README file there.

The script run.py is responsible for validating the configs and initiating the operation mode. Therefore, to start training or inference, run the following line with your configs file

python scr/run.py configs/<configs_file>.json

If you incorrectly defined any configs, an error will be thrown. Otherwise, operation will commence!

When you are finished, deactivate the environment

conda deactivate

3. Results

Results will be placed generated results/ directory with a unique directory name describing the run, or it can be given a specific name in the configs input file.

If training, each results directory will contain the following

• Copy of the input configs used
• .keras model file containing model structure and weights
• Model predictions for training and validation images
• CSV file with loss and metrics calculated during training
• Loss and metrics plots by epoch
• Results from each fold (cross validation only)
• Statistical plots of loss and metrics by epoch (cross validation only)

If inferencing, there will instead be

• Copy of the input configs used
• Model predictions

References

[1] O. Ronneberger, P. Fischer, and T. Brox, “U-Net: Convolutional Networks for Biomedical Image Segmentation,” arXiv.org, May 18, 2015. https://arxiv.org/abs/1505.04597
[2] Z. Zhou, M. Rahman, N. Tajbakhsh, and J. Liang, “UNet++: A Nested U-Net Architecture for Medical Image Segmentation,” arXiv.org, 2018. https://arxiv.org/abs/1807.10165
[3] M. Tan and Q. V. Le, “EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks,” arXiv.org, 2019. https://arxiv.org/abs/1905.11946