This is a PyTorch implementation of training and evaluation of U-Net deep learning model for semantic segmentation of materials microscopic images. The U-Net model proves to be able to achieve good performance for various materials image datasets. Some highlights of the implementation are
- Consistent format for semantic segmentation dataset
- Hyperparameter configuration
- Data augmentation visualization
The PyTorch and other necessary packages can be installed with pip. To do so,
run
pip install -r requirements.txt
The dataset folder need to be arranged in the following form:
data
|--DATASET
|--images
a.tif (or a.png, a.jpg)
|--labels
a.npy
|--labels_colored
a.png
|--splits
split_cv.csv
train.txt
The labels folder contains labels in numpy array form with data type
np.int8. The shape of the array should be two dimension: (width, height) and
each entry in the array gives the class of the pixel. Users can use image
reading tools such as pillow or matplotlib to load image into a numpy array and
use np.save to save it in .npy form. This process is to make sure that
labels take the same format for different dataset considering the diverse
original file format of labels. Note that the name of the label numpy file is
the same as the name of the image (a.tif, a.npy). The images in
labels_colored folder show the labels in image format and are used only for
visualization.
The files inside splits folder are used to split the dataset into train,
validation and test set for model development. There are two format of split files:
.csv: Containsnameandsplitcolumns. Thesplitcolumn specifies the data split index the image belongs. This split format is especially designed for cross-validation. See./data/uhcs/splits/split_cv.csvfor an example..text: Contains a list of image names in the dataset split. See./data/uhcs/splits/train16A.txtfor an example.
We use yaml files to store configurations for datasets and model development.
A default.yaml file specifying the default configuration should be created
under ./configs/DATASET for a new segmentation dataset. As an example,
the contents of ./configs/uhcs/default.yaml are
# dataset basics
dataset: uhcs # dataset name
img_folder: images # folder containing the images
label_folder: labels # folder containing the labels
n_classes: 4 # number of classes
# train, validation and test split
split_info:
type: CSVSplit # type of dataset split (CSVSplit, TextSplit)
test_type: CSVSplit # type of test split (CSVSplit, TextSplit, validation, folder)
split_file: split_cv.csv # name of split file
split_col_name: split # split column name of the csv split file
val_split_num: 1 # split index for validation set (csv split only)
test_split_num: 0 # split index for test set (csv split only)
train_reverse: True # if training set contains the remaining images
cross_validation:
val_splits: [1, 2, 3, 4, 5, 0] # validation split indices for k-fold cross-validation
test_splits: [0, 1, 2, 3, 4, 5] # test split indices for k-fold cross-validation
# data information
train_size: [224, 224] # size of cropped training images
eval_size: [484, 645] # size of evaluation images
mean: [0.485, 0.456, 0.406] # mean value in data normalization (ImageNet mean here)
std: [0.229, 0.224, 0.225] # standard deviation in data normalization (ImageNet std here)
# data augmentation
augmentations: # each line is a (operation: parameters) pair in the Albumentation package
Flip: {p: 0.5}
ShiftScaleRotate: {shift_limit: 0, scale_limit: [0,1], rotate_limit: 45, p: 0.8}
RandomBrightnessContrast: {brightness_limit: 0.2, contrast_limit: 0.3, p: 0.8}
# training
n_epochs: 150 # number of epochs
train_repeat: 8 # number of times to repeat training set per epoch
loss_type: CE # type of loss function (CE, Dice, Jaccard)
ignore_index: -1 # index to ignore in labels for computing loss
batch_size: 4 # size of training batches
optimizer:
type: AdamW # type of optimizer (AdamW, Adam, SGD)
encoder_lr: 5.e-5 # learning rate for U-
6BFA
Net encoder
decoder_lr: 5.e-4 # learning rate for U-Net decoder
weight_decay: 1.e-2 # weight decay (L2 penalty)
lr_scheduler:
type: MultiStepLR # type of learning rate Scheduler (MultiStepLR, CAWR, CyclicLR, OneCycleLR)
params: {milestones: [100], gamma: 0.3} # parameters of the learning rate Scheduler
metric: mIoU # validation metric to optimize for model selection (mIoU, IoU_i)
To run experiments with customized configurations, create a CONFIG.yaml
file specifying the different configuration parameters. For example,
example_textsplit.yaml and example_testfolder.yaml under ./configs/uhcs
give two other common scenarios for dataset splits.
To train a model, run train.py with the following argument:
python train.py --dataset DATASET --config CONFIG.yaml --gpu_id GPU_ID --seed SEED
dataset: name of the dataset (default:uhcs)config: configuration file (default:default.yaml)gpu_id: which GPU to run the script, cpu will be used if GPU is unavailable (default:0)seed: random seed to reproduce experiments (default:42)
During training, the model with the best metric score on validation set will be
saved in ./checkpoints/DATASET/CONFIG/model.pth.
To evaluation the model trained by CONFIG.yaml on validation and test set, run
eval.py with the following argument:
python eval.py --dataset DATASET --config CONFIG.yaml --gpu_id GPU_ID --mode MODE
- mode:
valortestspecifying validation or test
Evaluation will save the predictions in ./checkpoints/DATASET/CONFIG/predictions.
To train and evaluate models using GPU resources on Google Colab, please go to the notebook
colab.ipynb.
DeCost, B., Lei, B., Francis, T., & Holm, E. (2019). High Throughput Quantitative Metallography for Complex Microstructures Using Deep Learning: A Case Study in Ultrahigh Carbon Steel. Microscopy and Microanalysis, 25(1), 21-29. https://doi.org/10.1017/S1431927618015635 arxiv: https://arxiv.org/abs/1805.08693
Durmaz, A.R., Müller, M., Lei, B. et al. A deep learning approach for complex microstructure inference. Nat Commun 12, 6272 (2021). https://doi.org/10.1038/s41467-021-26565-5 arxiv: