By Qysoraka.

This repository consists of official PyTorch implementations of the papers:

• "Superpixel inpainting for self-supervised skin lesion segmentation from dermoscopic images", ISBI 2022
• "autoSMIM: Automatic Superpixel-based Masked Image Modeling for Skin Lesion Segmentation", TMI 2023

The code was developed using Python 3.8.16 and Pytorch 2.0.0 on CentOS 7 with NVIDIA GPUs. Training and testing were performed using 4 NVIDIA RTX 3090 GPUs with CUDA 11.7. Other platforms are not fully tested.

Main dependencies: Pytorch, Pytorch Lightning, WandB, Microsoft NNI. You can directly install the dependencies by running:

pip install -r requirements.txt

Note: If you want to use WandB to log your experiment results, you should have a free account of WandB. The experiment results will be automatically uploaded. For more information please refer to section below.

Please download ISIC datasets and Ham10000 datasets. After downloading the datasets, please change the function load_name in corresponding file in directory datasets to your data storage path (annotated with 'need modification').

Please refer to the configuration file in directory configs. All the training configurations can be done by changing the specific parameter in config file.

There are five config files in total.

• config.yaml: This config is used for original segmentation training and testing.
• config_inpainting.yaml: This config is used for MIM training and testing. Note: In this config file, parameter inpainting should be always set to be True. You should also change the parameters aug_k and aug_n to a specific value in the policy.
• config_search.yaml: This config is used for policy searching. Note: In this config file, parameter optimization should be always set to be True.
• nni_config.yaml: This config is used for nni setting. Here you can change the parameters of Bayesian Optimization experiments. I recommend only changing experimentName, trailConcurrency, maxTrialNumber and maxExperimentDuration to your own setting. For more information about these parameters, please follow the official website of Microsoft NNI.
• space.json: This config is used for policy search space setting. By default, the searching space is $[40, 1000]$ for $k$ and $[1, 500]$ for $n$. You can also specify another searching space for more experiments.

After successfully preparing your code, environment, and dataset, you can start your training. The training could be separated into four steps:

• Training MIM, rotation, colorization and jigsaw model using the base policy.

  You can do this training by running:

  python run.py --cfg ./configs/config_inpainting.yaml
  

  The policy is controlled by the parameters aug_k and aug_n. By default, the value of aug_k and aug_n is $[512, 128]$. You can change the value to your own.

• Policy optimization.

  With the config file set correctly, this step is simply done by running:

  nnictl create --config ./configs/nni_config.yaml --port 8080
  

  The policy optimization is implemented using Microsoft NNI. For customized usage of nni, please refer to its official repo.

• Training MIM model using optimized policy obtained in step 2.

  This is the same as the process in step 1. The only change is that you should change the value of aug_k, aug_n to an optimized policy value in step 2. The command is the same as in step 1.

• Downstream segmentation model training.

  After obtaining a new MIM model using an optimized policy, you can use the MIM model for segmentation. By setting the pretrained value in the config file config.yaml, you can directly load the MIM model for training segmentation model. Then directly run

  python run.py --cfg ./configs/config.yaml
  

When training the MIM model, we will need to generate superpixels for each image. This process is done automatically when loading the image in the data loader. However, this will slow down the whole process since it needs to generate a superpixel map for each image repeatedly.

If you want to speed up the training, you can generate a superpixel map before training and store the map in the directory of the corresponding dataset. You can do this by running sp_generation.py in the directory datasets.

python sp_generation.py

Here you should change the parameter corneal_list to your dataset storage path and change the parameter sp_num and size to specific data size and superpixels generated for each image. After doing so, the superpixel map will be automatically loaded from the directory of the corresponding dataset when loading the dataset.

The results of the experiments could be automatically uploaded to WandB. If you want to have the results uploaded, you need to create a free account of WandB. Please refer to their official website.

After getting a free account and signing in successfully, you can change the information in run.py to your own to log results automatically. If you don't want to use WandB, you can directly change the value of smoke_test in the config file to True.

If you find this repository useful, please consider citing the following papers:

@ARTICLE{10168139,
  author={Wang, Zhonghua and Lyu, Junyan and Tang, Xiaoying},
  journal={IEEE Transactions on Medical Imaging},
  title={autoSMIM: Automatic Superpixel-based Masked Image Modeling for Skin Lesion Segmentation},
  year={2023},
  volume={},
  number={},
  pages={1-1},
  doi={10.1109/TMI.2023.3290700}}

@inproceedings{wang2022superpixel,
  title={Superpixel inpainting for self-supervised skin lesion segmentation from dermoscopic images},
  author={Wang, Zhonghua and Lyu, Junyan and Luo, Wenhao and Tang, Xiaoying},
  booktitle={2022 IEEE 19th International Symposium on Biomedical Imaging (ISBI)},
  pages={1--4},
  year={2022},
  organization={IEEE}
}