Package: mist-medical
Medical Imaging Segmentation ToolkitOverviewThe Medical Imaging Segmentation Toolkit (MIST) is a simple, fully automated 3D medical imaging segmentation
framework. MIST allows researchers to quickly set up, train, and test a variety of deep learning models for 3D
medical imaging segmentation. The following architectures are implemented on MIST:nnUNetU-NetFMG-NetW-NetAttention U-NetUNETRThe following features are supported by MIST:NVIDIA Data Loading Library (DALI)Automatic mixed precision (AMP)Multi-GPU training with DistributedDataParallelPlease cite the following papers if you use this code for your work:A. Celaya et al., "PocketNet: A Smaller Neural Network For Medical Image Analysis," in IEEE Transactions on
Medical Imaging, doi: 10.1109/TMI.2022.3224873.A. Celaya et al, "FMG-Net and W-Net: Multigrid Inspired Deep Learning Architectures For Medical Imaging Segmentation", in
Proceedings of LatinX in AI (LXAI) Research Workshop @ NeurIPS 2023, doi: 10.52591/lxai202312104What's NewFeb. 2024 - MIST is now available as PyPI package and as a Docker image on DockerHub.Feb. 2024 - Major improvements to the analysis, preprocessing, and postprocessing pipelines,
and new network architectures like UNETR added.Feb. 2024 - We have moved the TensorFlow version of MIST tomist-tf.Table of ContentsOverviewSetupInstallData FormatGetting StartedOutputAdvanced UsageSetupMIST assumes that your system as at least one GPU and sufficient memory to handle 3D medical images.InstallTo install the latest version of MIST as an out-of-the-box segmentation pipeline, usepip install mist-medicalIf you want to install MIST and customize the underlying code (i.e., add a loss function or new architecture),
then clone the MIST repo and install as follows:git clone https://github.com/aecelaya/MIST.git
cd MIST
pip install -e .Data FormatThe MIST pipeline assumes that your train and test data directories are set up in the following structure.data/
│   └── patient_1
│       └── image_1.nii.gz
│       └── image_2.nii.gz
│       ...
│       └── image_n.nii.gz
|       └── mask.nii.gz
|
│   └── patient_2
│       └── image_1.nii.gz
│       └── image_2.nii.gz
│       ...
│       └── image_n.nii.gz
|       └── mask.nii.gz
|
|   ...
└──Please note that the naming convention here is for this example only. There is no specific
naming convention for the files within your dataset. However, the MIST pipeline assumes that the naming of the
images in each patient directory is consistent or that each type of image is identifiable by a
list of identifier strings.MIST offers support for MSD and CSV formatted datasets. For more details, please seeMSD and CSV Formatted Data.Once your dataset is in the correct format, the final step is to prepare a small JSON
file containing the details of the dataset. We specifically ask for the following key-value
pairs.KeyValuetaskName of task (i.e., brats, lits, etc.).modalityOptions arect,mr, orother.train-dataFull path to training data.test-dataFull path to test data (optional).maskList containing identifying strings for mask or ground truth images in dataset.imagesDictionary where each key is an image type (i.e., T1, T2, CT, etc.) and each value is a list containing identifying strings for that image type.labelsList of labels in dataset (starting with 0).final_classesDictionary where each key is the name of the final segmentation class (i.e., WT, ET, TC for BraTS) and each value is a list of the labels in that class.Here is an example for the BraTS 2023 dataset.{
    "task": "brats2023",
    "modality": "mr",
    "train-data": "/full/path/to/raw/data/train",
    "test-data": "/full/path/to/raw/data/validation",
    "mask": ["seg.nii.gz"],
    "images": {"t1": ["t1n.nii.gz"],
               "t2": ["t2w.nii.gz"],
               "tc": ["t1c.nii.gz"],
               "fl": ["t2f.nii.gz"]},
    "labels": [0, 1, 2, 3],
    "final_classes": {"WT": [1, 2, 3],
                      "TC": [1, 3],
                      "ET": [3]}
}Examples for several datasets are provided in theexamples/dataset-jsondirectory.Getting StartedThe MIST pipeline consists of three stages:Analyze - Gathering parameters about the dataset like target spacing, normalization parameters, etc. This produces aconfig.jsonfile, which will be used for the rest of the pipeline.Preprocess - Use the learned parameters from the analysis phase of the MIST pipeline to preprocess the data (i.e., reorient, resample, etc.) and convert it to numpy files.Train/Postprocess - Train on preprocessed data using a five-fold cross validation to produce a final set of models for inference. After training, a postprocessing analysis pipeline is activated to determine optimal postprocessing procedures.Additionally, MIST provides auxiliary commands that handle test-time prediction, evaluating a given set of predictions, and converting other datasets to the MIST format.When you install the MIST package, the following commands are included:mist_run_all: This command runs the entire MIST pipeline and requires the following arguments:--data: The full path to your dataset JSON file--numpy: The full path to the directory to save preprocessed Numpy files--results: The full path to the directory to save the output of the MIST pipeline--amp: Optional, but highly recommended, if your system supports AMP--pocket: Optional, but highly recommended, for using smaller, but just as accurate networksmist_analyze: This command runs only the analysis portion of the pipeline and requires the following arguments:--data: The full path to your dataset JSON file--results: The full path to the directory to save the output of the MIST pipelinemist_preprocess: This command runs only the preprocessing pipeline and assumes that the output of the analysis pipeline (or a modified version of it) are in the--resultsfolder. This command requires the following arguments:--data: The full path to your dataset JSON file--numpy: The full path to the directory to save preprocessed Numpy files--results: The full path to the directory to save the output of the MIST pipelinemist_train: This command runs only the training pipeline and assumes that the results of the previous two commands are in the--resultsand--numpyfolders. This command requires the following arguments:--data: The full path to your dataset JSON file--numpy: The full path to the directory to save preprocessed Numpy files--results: The full path to the directory to save the output of the MIST pipeline--amp: Optional, but highly recommended, if your system supports AMP--pocket: Optional, but highly recommended, for using smaller, but just as accurate networksmist_predict: This command runs test time inference on a given set of test data given as either a JSON or CSV file. This command is a stand alone command that can be used anytime outside of the MIST pipeline. To run this command, we need the following arguments:--models: The full path to themodelsfolder in the output of the MIST pipeline--config: The full path to the directory to theconfig.jsonfile in the output of the MIST pipeline--data: The full path to the JSON or CSV file, which contains the path to your test data--output: The full path to the directory to save the predictionsFor CSV formatted data, the CSV file must, at a minimum, have anidcolumn with the new patient IDs and one column for each image type. A column for themaskis allowed if you want to run the evaluation portion of the pipeline. For example, for the BraTS dataset, our CSV header would look like the following.idmask (optional)t1t2tcflPatient IDPath to maskPath to t1 imagePath to t2 imagePath to tc imagePath to fl imageSimilarly, for JSON formatted data, we would have the following.{
    "Patient ID": {
        "mask": "Path to mask", *** (optional) ***
        "t1": "Path to t1 image",
        "t2": "Path to t2 image",
        "tc": "Path to tc image", 
        "fl": "Path to fl image"
    }
}mist_evaluate: This command is used to evaluate a set of predictions and requires the following arguments:--paths: The full paths to the ground truth masks formated in either CSV or JSON format (seemist_predict)--config: The full path to the directory to theconfig.jsonfile in the output of the MIST pipeline--preds-dir: The full path to the directory where the predictions are saved--output-csv: The full path to the CSV file where the results will be savedmist_convert_dataset: This command converts either MSD or CSV formated datasets into MIST formatted data. The following arguments are required for this command:--format: The format of the given dataset, which can be eithermsdorcsv--msd-source: The full path to the MSD dataset, if that is what you are converting--csv-source: The full path to the CSV dataset, if that is what you are using--dest: The full path to the new MIST formatted datasetNote that the CSV file should mirror the CSV format shown in themist_predictcommand. Additionally, this command will reformat the CSV dataset to a MIST-compatible one but will require the user to fill in details in its corresponding dataset JSON file.DockerExcept formist_convert_dataset, each of the commands described above is also available as a separate Docker application. Start by pulling the Docker image you need with the following command:docker pull mistmedical/mist_run_all:latestTo run this, use the following command:docker run --rm -it -u $(id -u):$(id -g) --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -v /your/working/directory/:/workspace mist_run_all [--your arguments]OutputThe output of the MIST pipeline has the following structure:results/
│   └── logs/
│   └── models/
│   └── predictions/
│   config.json
│   results.csv
│   train_paths.csv
│   test_paths.csv (if test set is specified in dataset.json)
│   fg_bboxes.csv
└──Here is a breakdown of what each file/directory contains.Directory/FileDescriptionlogs/TensorBoard logs for each fold.models/PyTorch models for each fold andmodel_config.json, which stores the settings (i.e., architecture, patch size) to load the model from the saved weights.predictions/Raw and postprocessed predictions from five fold cross validation and test set (if specified).config.jsonJSON file containing all of the attributes of the dataset (i.e., target spacing, crop to foreground, etc.).results.csvCSV file with the Dice, 95th percentile Hausdorff, and average surface distance scores for all of the folds.train_paths.csvCSV file starting with columnsidandfoldspecifying the patient ID and which fold they belong to, and paths to the mask and images.test_paths.csvSame astrain_paths.csv, but for the test set if it's given.fg_bboxes.csvCSV file containing information about the bounding box around the foreground for each image.Advanced UsageAll MIST commands come with--helpor-hoption, which allows you to see all the available settings/arguments for that command.For themist_run_all,mist_analyze,mist_preprocess, andmist_traincommands, here is a complete list of the available arguments:usage: mist_run_all [-h] [--exec-mode {all,analyze,preprocess,train}]
                    [--data DATA] [--gpus GPUS [GPUS ...]]
                    [--num-workers NUM_WORKERS] [--master-port MASTER_PORT]
                    [--seed_val SEED_VAL] [--tta [BOOLEAN]]
                    [--results RESULTS] [--numpy NUMPY] [--amp [BOOLEAN]]
                    [--batch-size BATCH_SIZE]
                    [--patch-size PATCH_SIZE [PATCH_SIZE ...]]
                    [--max-patch-size MAX_PATCH_SIZE [MAX_PATCH_SIZE ...]]
                    [--learning-rate LEARNING_RATE] [--exp_decay EXP_DECAY]
                    [--lr-scheduler {constant,cosine_warm_restarts,exponential}]
                    [--cosine-first-steps COSINE_FIRST_STEPS]
                    [--optimizer {sgd,adam,adamw}] [--clip-norm [BOOLEAN]]
                    [--clip-norm-max CLIP_NORM_MAX]
                    [--model {nnunet,unet,attn_unet,unetr}]
                    [--use-res-block [BOOLEAN]] [--pocket [BOOLEAN]]
                    [--depth DEPTH] [--init-filters INIT_FILTERS]
                    [--deep-supervision [BOOLEAN]]
                    [--deep-supervision-heads DEEP_SUPERVISION_HEADS]
                    [--vae-reg [BOOLEAN]] [--vae-penalty VAE_PENALTY]
                    [--l2-reg [BOOLEAN]] [--l2-penalty L2_PENALTY]
                    [--l1-reg [BOOLEAN]] [--l1-penalty L1_PENALTY]
                    [--oversampling OVERSAMPLING] [--no-preprocess [BOOLEAN]]
                    [--use-n4-bias-correction [BOOLEAN]]
                    [--use-config-class-weights [BOOLEAN]]
                    [--class-weights CLASS_WEIGHTS [CLASS_WEIGHTS ...]]
                    [--loss {dice_ce,dice,gdl,gdl_ce}]
                    [--sw-overlap SW_OVERLAP]
                    [--blend-mode {constant,gaussian}]
                    [--no-postprocess [BOOLEAN]] [--nfolds NFOLDS]
                    [--folds FOLDS [FOLDS ...]] [--epochs EPOCHS]
                    [--steps-per-epoch STEPS_PER_EPOCH]
                    [--output-std [BOOLEAN]]

optional arguments:
  -h, --help            show this help message and exit
  --exec-mode {all,analyze,preprocess,train}
                        Run all of the MIST pipeline or an individual
                        component (default: all)
  --data DATA           Path to dataset json file (default: None)
  --gpus GPUS [GPUS ...]
                        Which gpu(s) to use, defaults to all available GPUs
                        (default: [-1])
  --num-workers NUM_WORKERS
                        Number of workers to use for data loading (default: 8)
  --master-port MASTER_PORT
                        Master port for multi-gpu training (default: 12355)
  --seed_val SEED_VAL   Random seed (default: 42)
  --tta [BOOLEAN]       Enable test time augmentation (default: False)
  --results RESULTS     Path to output of MIST pipeline (default: None)
  --numpy NUMPY         Path to save preprocessed numpy data (default: None)
  --amp [BOOLEAN]       Enable automatic mixed precision (recommended)
                        (default: False)
  --batch-size BATCH_SIZE
                        Batch size (default: 2)
  --patch-size PATCH_SIZE [PATCH_SIZE ...]
                        Height, width, and depth of patch size (default: None)
  --max-patch-size MAX_PATCH_SIZE [MAX_PATCH_SIZE ...]
                        Max patch size (default: [256, 256, 128])
  --learning-rate LEARNING_RATE
                        Learning rate (default: 0.001)
  --exp_decay EXP_DECAY
                        Exponential decay factor (default: 0.9)
  --lr-scheduler {constant,cosine_warm_restarts,exponential}
                        Learning rate scheduler (default: constant)
  --cosine-first-steps COSINE_FIRST_STEPS
                        Length of a cosine decay cycle in steps, only with
                        cosine_annealing scheduler (default: 500)
  --optimizer {sgd,adam,adamw}
                        Optimizer (default: adam)
  --clip-norm [BOOLEAN]
                        Use gradient clipping (default: False)
  --clip-norm-max CLIP_NORM_MAX
                        Max threshold for global norm clipping (default: 1.0)
  --model {nnunet,unet,attn_unet,unetr}
  --use-res-block [BOOLEAN]
                        Use residual blocks for nnUNet or UNet (default:
                        False)
  --pocket [BOOLEAN]    Use pocket version of network (default: False)
  --depth DEPTH         Depth of U-Net or similar architecture (default: None)
  --init-filters INIT_FILTERS
                        Number of filters to start network (default: 32)
  --deep-supervision [BOOLEAN]
                        Use deep supervision (default: False)
  --deep-supervision-heads DEEP_SUPERVISION_HEADS
                        Number of deep supervision heads (default: 2)
  --vae-reg [BOOLEAN]   Use VAE regularization (default: False)
  --vae-penalty VAE_PENALTY
                        Weight for VAE regularization loss (default: 0.1)
  --l2-reg [BOOLEAN]    Use L2 regularization (default: False)
  --l2-penalty L2_PENALTY
                        L2 penalty (default: 1e-05)
  --l1-reg [BOOLEAN]    Use L1 regularization (default: False)
  --l1-penalty L1_PENALTY
                        L1 penalty (default: 1e-05)
  --oversampling OVERSAMPLING
                        Probability of crop centered on foreground voxel
                        (default: 0.4)
  --no-preprocess [BOOLEAN]
                        Turn off preprocessing (default: False)
  --use-n4-bias-correction [BOOLEAN]
                        Use N4 bias field correction (only for MR images)
                        (default: False)
  --use-config-class-weights [BOOLEAN]
                        Use class weights in config file (default: False)
  --class-weights CLASS_WEIGHTS [CLASS_WEIGHTS ...]
                        Specify class weights (default: None)
  --loss {dice_ce,dice,gdl,gdl_ce}
                        Loss function for training (default: dice_ce)
  --sw-overlap SW_OVERLAP
                        Amount of overlap between scans during sliding window
                        inference (default: 0.25)
  --blend-mode {constant,gaussian}
                        How to blend output of overlapping windows (default:
                        constant)
  --no-postprocess [BOOLEAN]
                        Run post processing on MIST output (default: False)
  --nfolds NFOLDS       Number of cross-validation folds (default: 5)
  --folds FOLDS [FOLDS ...]
                        Which folds to run (default: [0, 1, 2, 3, 4])
  --epochs EPOCHS       Number of epochs (default: 300)
  --steps-per-epoch STEPS_PER_EPOCH
                        Steps per epoch. By default ceil(training_dataset_size
                        / batch_size / gpus) (default: None)
  --output-std [BOOLEAN]
                        Output standard deviation for ensemble predictions
                        (default: False)Here are the available arguments formist_predict:usage: mist_predict [-h] [--models MODELS] [--config CONFIG] [--data DATA]
                    [--output OUTPUT] [--fast [BOOLEAN]] [--gpu GPU]
                    [--sw-overlap SW_OVERLAP]
                    [--blend-mode {constant,gaussian}] [--tta [BOOLEAN]]
                    [--no-preprocess [BOOLEAN]] [--output_std [BOOLEAN]]

optional arguments:
  -h, --help            show this help message and exit
  --models MODELS       Directory containing saved models (default: None)
  --config CONFIG       Path and name of config.json file from results of MIST
                        pipeline (default: None)
  --data DATA           CSV or JSON file containing paths to data (default:
                        None)
  --output OUTPUT       Directory to save predictions (default: None)
  --fast [BOOLEAN]      Use only one model for prediction to speed up
                        inference time (default: False)
  --gpu GPU             GPU id to run inference on (default: 0)
  --sw-overlap SW_OVERLAP
                        Amount of overlap between scans during sliding window
                        inference (default: 0.25)
  --blend-mode {constant,gaussian}
                        How to blend output of overlapping windows (default:
                        constant)
  --tta [BOOLEAN]       Use test time augmentation (default: False)
  --no-preprocess [BOOLEAN]
                        Turn off preprocessing (default: False)
  --output_std [BOOLEAN]
                        Outputs standard deviation image (default: False)Here are the available arguments formist_evaluate:usage: mist_evaluate [-h] [--config CONFIG] [--paths PATHS]
                     [--preds-dir PREDS_DIR] [--output-csv OUTPUT_CSV]

optional arguments:
  -h, --help            show this help message and exit
  --config CONFIG       Path to config.json file from MIST output (default:
                        None)
  --paths PATHS         Path to CSV or JSON file with original mask/data
                        (default: None)
  --preds-dir PREDS_DIR
                        Path to directory containing predictions (default:
                        None)
  --output-csv OUTPUT_CSV
                        Path to CSV containing evaluation results (default:
                        None)Here are the available arguments formist_convert_dataset:usage: mist_convert_dataset [-h] [--format {msd,csv}]
                            [--msd-source MSD_SOURCE] [--train-csv TRAIN_CSV]
                            [--test-csv TEST_CSV] [--dest DEST]

optional arguments:
  -h, --help            show this help message and exit
  --format {msd,csv}    Format of dataset to be converted (default: msd)
  --msd-source MSD_SOURCE
                        Directory containing MSD formatted dataset (default:
                        None)
  --train-csv TRAIN_CSV
                        Path to and name of csv containing training ids, mask,
                        and images (default: None)
  --test-csv TEST_CSV   Path to and name of csv containing test ids and images
                        (default: None)
  --dest DEST           Directory to save converted, MIST formatted dataset
                        (default: None)