Unet is an architecture used widely in medical field to solve the high resolution problem on most of the medical images. There are lots of version of the Unet architecture. However, it is the pytorch version or the old keras version, and thus can't get the advantage of the new vesrion tensorflow.Therefore, I provide a tensorflow 2.2.0 version and the simple architecture code for people to use and modify.
1.MakeDataset.py
For the purpose of making a dataset by the tensorflow iterator, tf.data.Dataset, which is faster than
using python list to read images.
2.Unet_model.py
To make the Unet model by using tensorflow inheritance, tf.keras.Model. You can use the Model func for
sure, it's up to you.
3.utils.py
Provide the utilities for you to visualize the results of your AI model. There are also some image
processing methods for you to enhance the segmentation results using the opencv.
4.Unet_main.py
The main function to control the whole process of the Unet model training. There are lots of thing you
can tune yourself in it, like the hyperparameters and the callbacks. Feel free to modify it yourself.
Be sure to check all the steps below to make sure nothing goes wrong in your training process.
- Data leakage -> divide the dataset by patients
- Different dataset distribution -> use five fold method
- Save all the train values -> don't miss any train/valid values or the initial training hyperparameters
- Plot the result for you to check -> check whether the model got the right features or not
Unet
├── Train
│ ├── images
│ └── masks
└── Validation
├── images
└── masks
find $PWD/Train/images/* -iname "*.jpg" | awk -F'/' 'BEGIN{OFS=",";print"Path,Y,ID"}{print $0,$6,$8}' > Train_image.csv
Hint : Be sure to change the "$0,$6,$8" term. Modify them dependng on the path of your images. You can see the command awk for more details.
find $PWD/Train/masks/ -iname "*.jpg" | awk -F'/' 'BEGIN{OFS=",";print"Path,Y,ID"}{print $0,$6,$8}' > Train_mask.csv
Hint:You can make yor csv file in any form, just make sure you have the three columns:["Path","Y","ID"]
Assume you have all the needed dataset. If not, go through the Pre-requisites and get all the needed datasets.
python Unet.main.py
# You can modify the hyperparameters in the Unet_main.py, like:
# epoch = 1000
# learning rate = 0.0001
# callbacks = [x, y, z]
# metrics = IOU
# optimizer = adam
Hint : Make sure to change all the file paths of yours, and also be aware of the different naming method between mine and yours.