This repository contains an implementation of the convolutional neural network (CNN) described in this paper. The CNN segments biomedical images by predicting the class of each pixel. In the original task, the images were of cardiac histological sections stained with Masson's trichrome. The objective was to segment the RGB images into 3 classes: myocyte, background, or fibrosis.
The code requires Python and TensorFlow to run.
Training data should be in the ./training set
folder. RGB images should be named in the format train_1.png
, train_2.png
, etc. Ground truth masks should be named in the format train_1_mask.png
, train_2_mask.png
, etc.
In train.py
, ensure the variables in the list from batch_size
to keep_rate
are correct.
If restore
is set to True
, a previous model will first be loaded and training will resume. Change the path of the model being restored under the function train_network
in network.py
.
To train the model, run train.py
A model is saved every epoch in the ./model
folder, TensorBoard logs are saved under ./logs
. One prediction is saved per epoch under ./predictions training
.
Test data should be in the ./testing set
folder. RGB images should be named in the format test_1.png
, test_2.png
, etc. If available, ground truth masks should be named in the format test_1_mask.png
, test_2_mask.png
, etc.
In test.py
, ensure the variables in the list from n_epochs
to n_predict
are correct.
To test, run test.py
Predicted segmentations are saved under ./predictions test
.
Accuracy and DSC scores are computed for each test image.