Semantic segmentation and image-to-image translation based on AI. This repo collects methods for pre-process data, style transfer and semantic segmentation.

Dependencies for specific versions of CUDA are listed in requirements/cu113.txt, requirements/cu111.txt, etc.. It is advisable to install the package in a virtual environment, e.g. using venv

cd /your/working/directory
python -m venv .venv

Activate it using e.g. source .venv/bin/activate on Linux/Mac and .venv\Scripts\activate.bat on Windows.

To install this repo (this will install all dependencies):

pip install git+https://github.com/dyollb/segmantic.git#egg=segmantic

Or in edit/dev mode

pip install -e git+https://github.com/dyollb/segmantic.git#egg=segmantic[dev]

On Windows installing torch with GPU support is slightly more involved. Make sure to first install torch matching the installed CUDA version first or use the requirements files, e.g. for CUDA 11.3

--find-links https://download.pytorch.org/whl/cu113/torch_stable.html
torch==1.11.0+cu113
torchvision==0.12.0+cu113

Run training:

segmantic-unet train --help
segmantic-unet train -i work/inputs/images -l work/inputs/labels -t work/inputs/labels.txt -r work/outputs

Or with a config file - first create empty config file (yml or json format):

segmantic-unet train-config -c config.yml --print-defaults

Edit config.yml e.g. to

image_dir: work/inputs/images
labels_dir: work/inputs/labels
tissue_list: work/inputs/labels.txt
output_dir: work/outputs
checkpoint_file: null
num_channels: 1
spatial_dims: 3
max_epochs: 500
augment_intensity: true
augment_spatial: false
mixed_precision: true
cache_rate: 1.0
gpu_ids:
- 0

Now run training:

segmantic-unet train-config -c config.yml

The example above included a tissue_list option. This is a path to a text file specifying the labels contained in a segmented image. By convention the 'label=0' is the background and is omitted from the the format. A segmentation with three tissues 'Bone'=1, 'Fat'=2, and 'Skin'=3 would be specified as follows:

    V7
    N3
    C0.00 0.00 1.00 0.50 Bone
    C0.00 1.00 0.00 0.50 Fat
    C1.00 0.00 0.00 0.50 Skin

Instead of providing the 'image_dir'/'labels_dir' pair, the training data can also be described by one or multiple json files. Example config that globs data from multiple json files:

{
    "datalist": ["/dataA/dataset.json", "/dataB/dataset.json"],
    "output_dir": "<path where trained model and logs are saved>",
    "Etc": "etc"
}

The dataset.json loosely follows the convention used for the Medical Segmentation Decathlon datasets, and popular codes e.g. nnUNet.