The paper used for this project is: Skin Lesion Segmentation: U-Nets versus Clustering, written by Bill S. Lin, Kevin Michael, Shivam Kalra, H.R. Tizhoosh.
To use this project:
- Install the requirements:
pip install -r requirements.txt- Open the file preprocessing.py and follow the instructions in the docstring of the function build_training_set to create the data set (download images of ISIC challenge etc.). The data set is available here: https://challenge2018.isic-archive.com/task1/training/
- Run the file train.py to train a model. It will be saved in the saved_models folder.
- Run the file predict.py to predict some masks. The results will be displayed in the folder results/model_id. The model provided here has the id 2019-04-25_12-19-15.
- preprocessing.py
In this file, we transform the RGB images into 5-channels images as indicated in the paper. We just take a size of 320 x 320 instead of 342 x 342.
- data_generator.py
This file contains a class, DataGenerator, that allows us to manage the data set, which contains 5-channel images and binary masks.
- train.py
This file contains the code to train our model. We did try to fully implement the U-Net or to use the library tf_unet but we focused on the segmentation_models library which yielded better results. Besides, this model takes a size divisible by 32!
- predict.py
To predict a mask, see the last cell of the file to use your own images. The original images must be in the folder ISIC2018_Task1-2_Training_Input and the ground truth mask, if available, in the folder ISIC2018_Task1_Training_GroundTruth.
- utils.py
This file contains some functions to save a model, to plot some results etc.