Ecosyste.ms: Awesome
An open API service indexing awesome lists of open source software.
https://github.com/mateuszbuda/brain-segmentation-pytorch
U-Net implementation in PyTorch for FLAIR abnormality segmentation in brain MRI
https://github.com/mateuszbuda/brain-segmentation-pytorch
Last synced: 2 months ago
JSON representation
U-Net implementation in PyTorch for FLAIR abnormality segmentation in brain MRI
- Host: GitHub
- URL: https://github.com/mateuszbuda/brain-segmentation-pytorch
- Owner: mateuszbuda
- License: mit
- Created: 2019-05-26T00:27:21.000Z (about 5 years ago)
- Default Branch: master
- Last Pushed: 2023-03-24T22:54:48.000Z (about 1 year ago)
- Last Synced: 2024-01-26T20:35:46.529Z (5 months ago)
- Language: Python
- Homepage: https://mateuszbuda.github.io/2017/12/01/brainseg.html
- Size: 29.4 MB
- Stars: 663
- Watchers: 9
- Forks: 179
- Open Issues: 21
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Lists
- Awesome-pytorch-list - brain-segmentation-pytorch - Net implementation in PyTorch for FLAIR abnormality segmentation in brain MRI. (Paper implementations / Other libraries:)
- Awesome-pytorch-list-CNVersion - brain-segmentation-pytorch - Net的PyTorch模型实现,用于脑核磁共振中FLAIR异常的分割。 (Paper implementations|论文实现 / Other libraries|其他库:)
README
# U-Net for brain segmentation
U-Net implementation in PyTorch for FLAIR abnormality segmentation in brain MRI based on a deep learning segmentation algorithm used in [Association of genomic subtypes of lower-grade gliomas with shape features automatically extracted by a deep learning algorithm](https://doi.org/10.1016/j.compbiomed.2019.05.002).
This repository is an all Python port of official MATLAB/Keras implementation in [brain-segmentation](https://github.com/mateuszbuda/brain-segmentation).
Weights for trained models are provided and can be used for inference or fine-tuning on a different dataset.
If you use code or weights shared in this repository, please consider citing:```
@article{buda2019association,
title={Association of genomic subtypes of lower-grade gliomas with shape features automatically extracted by a deep learning algorithm},
author={Buda, Mateusz and Saha, Ashirbani and Mazurowski, Maciej A},
journal={Computers in Biology and Medicine},
volume={109},
year={2019},
publisher={Elsevier},
doi={10.1016/j.compbiomed.2019.05.002}
}
```## docker
```
docker build -t brainseg .
``````
nvidia-docker run --rm --shm-size 8G -it -v `pwd`:/workspace brainseg
```## PyTorch Hub
Loading model using PyTorch Hub: [pytorch.org/hub/mateuszbuda\_brain-segmentation-pytorch\_unet](https://pytorch.org/hub/mateuszbuda_brain-segmentation-pytorch_unet/)
```python
import torch
model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet',
in_channels=3, out_channels=1, init_features=32, pretrained=True)
```## data
![dataset](./assets/brain-mri-lgg.png)
Dataset used for development and evaluation was made publicly available on Kaggle: [kaggle.com/mateuszbuda/lgg-mri-segmentation](https://www.kaggle.com/mateuszbuda/lgg-mri-segmentation).
It contains MR images from [TCIA LGG collection](https://wiki.cancerimagingarchive.net/display/Public/TCGA-LGG) with segmentation masks approved by a board-certified radiologist at Duke University.## model
A segmentation model implemented in this repository is U-Net as described in [Association of genomic subtypes of lower-grade gliomas with shape features automatically extracted by a deep learning algorithm](https://doi.org/10.1016/j.compbiomed.2019.05.002) with added batch normalization.
![unet](./assets/unet.png)
## results
|![TCGA_DU_6404_19850629](./assets/TCGA_DU_6404_19850629.gif)|![TCGA_HT_7879_19981009](./assets/TCGA_HT_7879_19981009.gif)|![TCGA_CS_4944_20010208](./assets/TCGA_CS_4944_20010208.gif)|
|:-------:|:-------:|:-------:|
| 94% DSC | 91% DSC | 89% DSC |Qualitative results for validation cases from three different institutions with DSC of 94%, 91%, and 89%.
Green outlines correspond to ground truth and red to model predictions.
Images show FLAIR modality after preprocessing.![dsc](./assets/dsc.png)
Distribution of DSC for 10 randomly selected validation cases.
The red vertical line corresponds to mean DSC (91%) and the green one to median DSC (92%).
Results may be biased since model selection was based on the mean DSC on these validation cases.## inference
1. Download and extract the dataset from [Kaggle](https://www.kaggle.com/mateuszbuda/lgg-mri-segmentation).
2. Run docker container.
3. Run `inference.py` script with specified paths to weights and images. Trained weights for input images of size 256x256 are provided in `./weights/unet.pt` file. For more options and help run: `python3 inference.py --help`.## train
1. Download and extract the dataset from [Kaggle](https://www.kaggle.com/mateuszbuda/lgg-mri-segmentation).
2. Run docker container.
3. Run `train.py` script. Default path to images is `./kaggle_3m`. For more options and help run: `python3 train.py --help`.Training can be also run using Kaggle kernel shared together with the dataset: [kaggle.com/mateuszbuda/brain-segmentation-pytorch](https://www.kaggle.com/mateuszbuda/brain-segmentation-pytorch).
Due to memory limitations for Kaggle kernels, input images are of size 224x224 instead of 256x256.Running this code on a custom dataset would likely require adjustments in `dataset.py`.
Should you need help with this, just open an issue.## TensorRT inference
If you want to run the model inference with TensorRT runtime, here is a blog post from Nvidia that covers this: [Speeding Up Deep Learning Inference Using TensorRT](https://developer.nvidia.com/blog/speeding-up-deep-learning-inference-using-tensorrt/).