https://github.com/umessen/salt
Softmax for Arbitrary Label Trees (SALT) is a framework for training segmentation networks using conditional probabilities to model hierarchical relationships in the data.
https://github.com/umessen/salt
3d-segmentation conditional-probabilities hierarchical-labels medical-image-segmentation
Last synced: 5 months ago
JSON representation
Softmax for Arbitrary Label Trees (SALT) is a framework for training segmentation networks using conditional probabilities to model hierarchical relationships in the data.
- Host: GitHub
- URL: https://github.com/umessen/salt
- Owner: UMEssen
- License: mit
- Created: 2024-03-27T07:40:30.000Z (about 2 years ago)
- Default Branch: main
- Last Pushed: 2024-07-17T20:47:49.000Z (almost 2 years ago)
- Last Synced: 2024-07-18T01:04:18.685Z (almost 2 years ago)
- Topics: 3d-segmentation, conditional-probabilities, hierarchical-labels, medical-image-segmentation
- Language: Python
- Homepage: https://ship-ai.ikim.nrw/
- Size: 60.4 MB
- Stars: 73
- Watchers: 5
- Forks: 4
- Open Issues: 2
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
[](https://opensource.org/licenses/MIT)
# Softmax for Arbitrary Label Trees
This segmentation framework was developed at the Institute for AI in Medicine of the University Hospital Essen by the [SHIP.AI](https://ship-ai.ikim.nrw/) team.
The framework can be used for any 2D or 3D segmentation task that exhibits a hierarchical labels structure.
In our case, we applied this to medical imaging and provide the segmentation of 145 different structures in the human body.
The model is very fast (average of 35 seconds for a 1000 whole body!) and can segment 113 body regions in a single pass.
## Preprint & Paper
You can find the preprint of our paper on [here](https://arxiv.org/abs/2407.08878). If you use SALT please cite the tool as following:
```
Sven Koitka, Giulia Baldini, Cynthia S. Schmidt, Olivia B. Pollok, Obioma Pelka, Judith Kohnke, Katarzyna Borys, Christoph M. Friedrich, Benedikt M. Schaarschmidt, Michael Forsting, Lale Umutlu, Johannes Haubold, Felix Nensa, & René Hosch. (2024). SALT: Introducing a Framework for Hierarchical Segmentations in Medical Imaging using Softmax for Arbitrary Label Trees.
```
## Training
### Datasets
For training, the dataset folder should have the following structure:
```
data/
├── lits/
├── kits/
├── saros/
├── ...
```
Each dataset should contain the following files and folders:
```
data/
├── saros/
│ ├── labels.txt
| ├── tree-labels.txt
│ ├── train/
│ │ ├── images/
│ │ │ ├── s0001.nii.gz
│ │ │ ├── s0002.nii.gz
│ │ ├── labels/
│ │ │ ├── s0001.nii.gz
│ │ │ ├── s0002.nii.gz
│ ├── val/
│ │ ├── images
│ │ │ ├── ...
│ │ ├── labels
│ │ │ ├── ...
│ ├── test/
│ │ ├── images
│ │ │ ├── ...
│ │ ├── labels
```
The datasets do not need to have all three `train`, `val`, and `test` folders, and may only include a `test` set.
The files `labels.txt` and `tree-labels.txt` specify the labels from the dataset. `labels.txt` should contain the label names, while `tree-labels.txt` should define the hierarchical structure of the labels from the tree. In the [labels](./labels) folder, you can find some examples for these files, and in the [conversion](salt/data/conversion) you can find examples on how the data was converted to this format.
### Train
Once the data is set up, you can build the docker image with
```bash
docker build -t shipai/salt .
```
and then run the container with
```bash
docker run -it --rm \
--runtime=nvidia \
--network host \
--user $(id -u):$(id -g) \
--shm-size=8g --ulimit memlock=-1 --ulimit stack=67108864 \
-v /path/to/storage:/storage \
-e NVIDIA_VISIBLE_DEVICES=0,1,2 \
shipai/salt
```
Once the container is running, you can use the following commands to train the models:
```bash
python3 -m salt.train \
--data-dir /storage/datasets/data/ \
--mixed-precision \
--train-dir /storage/models/your-model-name \
--epochs 1000 --batch-size 3 --gpus 0 1 2 --sink-weight 0.0
```
The number of gpus and the batch size you use depends on how many GPUs you made available in your docker run command.
Once the training is finished, you can export the model to a single model file:
```bash
python -m salt.export \
--train-dir /storage/models/your-model-name \
--output-dir /storage/models/exports/your-model-name
```
## Prediction & Evaluation
You can build the docker image with
```bash
docker build -t shipai/salt .
```
and then run the container with
```bash
docker run -it --rm \
--runtime=nvidia \
--network host \
--user $(id -u):$(id -g) \
--shm-size=8g --ulimit memlock=-1 --ulimit stack=67108864 \
-v /path/to/storage:/storage \
-e NVIDIA_VISIBLE_DEVICES=0,1,2 \
shipai/salt
```
Once the container is running, you can compute the predictions for your dataset using our model
```bash
python -m salt.predict \
--data-dir /storage/datasets/data/dataset-name/images \
--output_dir /storage/results/your-model-name/dataset-name
```
or run it with your newly trained model
```bash
python -m salt.predict \
--model-file /storage/models/exports/your-model-name/model.pt \
--config-file /storage/models/exports/your-model-name/config.pkl \
--data-dir /storage/datasets/data/dataset-name/images \
--output_dir /storage/results/your-model-name/dataset-name
```
Note that for the predictions, the data does not need to be in a particular format, and any folder containing NIfTI files can be used as input.
If your dataset has a ground truth, you can evaluate the predictions from outside the container using:
```bash
poetry run python -m salt.classic_evaluate \
--config-file /path/to/the/model/config.pkl \
--data-dir /path/to/the/data \
--predictions-dir /path/to/the/predictions \
--output-dir /path/for/results
```