[MICCAI 2024] CUTS

A Deep Learning and Topological Framework for Multigranular Unsupervised Medical Image Segmentation

Krishnaswamy Lab, Yale University

This is the authors' PyTorch implementation of CUTS, MICCAI 2024.

The official version is maintained in the Lab GitHub repo.

A Glimpse into the Methods

Citation

@inproceedings{Liu_CUTS_MICCAI2024, title = { { CUTS: A Deep Learning and Topological Framework for Multigranular Unsupervised Medical Image Segmentation } }, author = { Liu, Chen and Amodio, Matthew and Shen, Liangbo L. and Gao, Feng and Avesta, Arman and Aneja, Sanjay and Wang, Jay C. and Del Priore, Lucian V. and Krishnaswamy, Smita}, booktitle = {proceedings of Medical Image Computing and Computer Assisted Intervention -- MICCAI 2024}, publisher = {Springer Nature Switzerland}, volume = {LNCS 15008}, page = {155–165}, year = {2024}, month = {October}, }

Repository Hierarchy

`` UnsupervisedMedicalSeg (CUTS) ├── (*) comparison: other SOTA unsupervised methods for comparison. | ├── checkpoints: model weights are saved here. ├── config: configuration yaml files. ├── data: folders containing data files. ├── logs: training log files. ├── results: generated results (images, labels, segmentations, figures, etc.). | └── src ├── (*) scripts_analysis: scripts for analysis and plotting. | ├──generatebaselines.py | ├──generatekmeans.py | ├──generatediffusion.py | ├──plotpaperfiguremain.py | └──runmetrics.py | ├── (*)main.py: unsupervised training of the CUTS encoder. ├── (*)mainsupervised.py: supervised training of UNet/nnUNet for comparison. | ├── datasets: defines how to access and process the data inCUTS/data/`. ├── data_utils ├── model └── utils

Relatively core files or folders are marked with (*). ```

Data Provided

As some background info, I inherited the datasets from a graduated member of the lab when I worked on this project. These datasets are already preprocessed by the time I had them. For reproducibility, I have included the berkeley_natural_images, brain_tumor and retina datasets in zip format in this directory. The brain_ventricles dataset exceeds the GitHub size limits, and can be found on Google Drive.

Please be mindful that these datasets are relatively small in sample size. If big sample size is a requirement, you can look into bigger datasets such as the BraTS challenge.

The entire dataset can also be downloaded from Huggingface.

To reproduce the results in the paper.

The following commands are using retina_seed2 as an example (retina dataset, random seed set to 2022).

To train on your custom dataset with label.

NOTE: Since our method is completely unsupervised, the only additional benefit for providing labels are: 1. With labels, we will have meaningful quantitative metrics (dice coefficient, etc.) when you run run_metrics.py. 2. With labels, for binary segmentation tasks, our method can predict binary masks in addition to multi-scale segmentations.

The process is largely the same as detailed in the section: To reproduce the results in the paper above.

Other than that, you can use the pipeline as usual.

To train on your custom dataset without label.

The process is largely the same as detailed in the section: To reproduce the results in the paper above.

Other than that, you can use the pipeline as usual.

Be mindful though: when you run generate_kmeans.py, the script will still print out dice scores for each image. The values shall be very close to zero (on the order of 1e-16). This does not mean the segmentation is bad. This only means the ground truth label is not provided. The dice score is computed against a placeholding all-zero label, with a very tiny numerical stability term.

SPECIAL NOTE: The outcome of this pipeline will be the multi-scale segmentations as a result of diffusion condensation. No binary mask will be generated. If you really want the model to generate binary masks in addition to the multi-scale segmentations, what you can do is to provide a set of pseudo-labels as follows: 1. Instead of segmenting the desired region-of-interest carefully as in regular labels, you just casually circle/square/whatever a typical enough subregion inside the region-of-interest, and mark them as 1's whereas the backgrounds as 0's. You can do this with any labeling tool you like. 2. Then, you provide these casual pseudo-labels as if they were real labels. Put them in the correct folder under src/data/ and load them as if they were real labels in your src/datasets/***.py. 3. Follow the pipeline as in the previous section: To train on your custom dataset with label. 4. In this way, binary masks will also be generated. Please be careful though: the quantitative metrics will not be accurate, as you are providing pseudo-labels instead of accurate real labels.

Dependencies

We developed the codebase in a miniconda environment. Tested on Python 3.9.13 + PyTorch 1.12.1. How we created the conda environment: ``` conda create --name cuts pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch conda activate cuts conda install scikit-image pillow matplotlib seaborn tqdm -c anaconda python -m pip install -U phate python -m pip install multiscale-phate python -m pip install git+https://github.com/KrishnaswamyLab/CATCH python -m pip install opencv-python-headless python -m pip install sewar python -m pip install monai python -m pip install nibabel

(Optional) For STEGO

python -m pip install omegaconf python -m pip install wget python -m pip install torchmetrics python -m pip install tensorboard python -m pip install pytorch-lightning==1.9 python -m pip install azureml python -m pip install azureml.core

(Optional) For SAM

python -m pip install git+https://github.com/facebookresearch/segment-anything.git

(Optional) For MedSAM

python -m pip install git+https://github.com/bowang-lab/MedSAM.git

(Optional) For SAM-Med2D

python -m pip install albumentations python -m pip install scikit-learn==1.1.3 # need to downgrade to 1.1.3 ``` Installation usually takes between 20 minutes and 1 hour on a normal desktop computer.

The installation method above has been verified on a few different devices. Alternatively, you can try to use the environment.yaml or requirements.txt to install, but these are not verified.

DEBUG Notes

Acknowledgements

For the comparison against other methods, we use the official implementations from the following repositories: - DFC, IEEE TIP 2020: Unsupervised Learning of Image Segmentation Based on Differentiable Feature Clustering - STEGO, ICLR 2022: Unsupervised Semantic Segmentation by Distilling Feature Correspondences - SAM, ICCV 2023 (Meta AI Research): Segment Anything - SAM, Medical Image Analysis 2024: Segment Anything Model for Medical Image Analysis: an Experimental Study - SAM-Med2D, ArXiv: SAM-Med2D - MedSAM, Nature Communications 2024: Segment anything in medical images