Medical Matting: A New Perspective on Medical Segmentation with Uncertainty, arxiv

This is a PyTorch implementation of our paper. We introduce matting as a soft segmentation method and a new perspective to deal with and represent uncertain regions into medical scenes.

Reference

@InProceedings{10.1007/978-3-030-87199-4_54,
author="Wang, Lin and Ju, Lie and Zhang, Donghao and Wang, Xin and He, Wanji and Huang, Yelin and Yang, Zhiwen and Yao, Xuan and Zhao, Xin and Ye, Xiufen and Ge, Zongyuan",
title="Medical Matting: A New Perspective on Medical Segmentation with Uncertainty",
booktitle="Medical Image Computing and Computer Assisted Intervention -- MICCAI 2021",
year="2021",
pages="573--583",
}

A jounal extension version can be achieved at arxiv.

1. Requirements

torch>=1.8.0; torchvision>=0.9.0; matplotlib; numpy; opencv_python; pandas; Pillow; PyYAML; scikit_image; scikit_learn; scipy; skimage; tensorboardX; tqdm; pickle;

data prepare: unpack the zip files in the dataset folder first.

Directory structure in this repo:

│MedicalMatting/
│   config.py
│   evaluate.py
│   params_brain.yaml
│   params_isic.yaml
│   params_lidc.yaml
│   README.md
│   train.py
+---dataloader
│       data_loader.py
│       data_spliter.py
│       transform.py
│       utils.py
+---dataset
│       brain_growth_alpha.pkl
│       isic_attributes.pkl
│       lidc_attributes.pkl
+---model
│   │   loss_functions.py
│   │   loss_strategy.py
│   │   medical_matting.py
│   │   utils.py
│   +---matting_network
│   │       cbam.py
│   │       matting_net.py
│   │       resnet_block.py
│   +---metrics
│   │       compute_connectivity_error.py
│   │       compute_gradient_loss.py
│   │       compute_mse_loss.py
│   │       compute_sad_loss.py
│   │       dice_accuracy.py
│   │       generalised_energy_distance.py
│   │       utils.py
│   +---probabilistic_unet
│   │       axis_aligned_conv_gaussian.py
│   │       encoder.py
│   │       fcomb.py
│   │       prob_unet.py
│   \---unet
│           unet.py
│           unet_blocks.py
+---models
\---utils
        logger.py
        utils.py

2. Train

LIDC-IDRI

CUDA_VISIBLE_DEVICES=0 python train.py --config /path/to/params_lidc.yaml

ISIC

CUDA_VISIBLE_DEVICES=0 python train.py --config /path/to/params_isic.yaml

Brain-growth

CUDA_VISIBLE_DEVICES=0 python train.py --config /path/to/params_brain.yaml

3. Evaluation

CUDA_VISIBLE_DEVICES=0. python evaluate.py --config /path/to/params_**task.yaml \
  --save_path /path/to/your/model/dir

Acknowledgements

The following code is referenced in this repo.

A PyTorch implementation of Probablistic UNet.
CBAM
We reimplement the evaluation metrics in Matlab DIM_evaluate_code with PyTorch, more details please refer to DIM_evaluation_code_python.

Datasets:

The datasets in this paper were constructed based on the LIDC-IDRI, ISIC, and Brain-growth dataset, and the rights to the images used are owned by the original datasets. Please refer to the requirements of the original datasets for any use of the original images.

LIDC-IDRI
- The authors acknowledge the National Cancer Institute and the Foundation for the National Institutes of Health, and their critical role in the creation of the free publicly available LIDC/IDRI Database used in this study.
- The LIDC-IDRI dataset used in this paper was obtained based on further annotation of these patches.
ISIC
Brain-growth

LICENSE

The codes of this repo is under the GPL license. For commercial use, please contact with the authors.