diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0d46c1d --- /dev/null +++ b/.gitignore @@ -0,0 +1,175 @@ +# Vscode +.vscode +.vscode/* +!.vscode/settings.json +!.vscode/tasks.json +!.vscode/launch.json +!.vscode/extensions.json +!.vscode/*.code-snippets + +# Local History for Visual Studio Code +.history/ + +# Built Visual Studio Code Extensions +*.vsix + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ diff --git a/Dockerfile b/Dockerfile deleted file mode 100644 index f2312dc..0000000 --- a/Dockerfile +++ /dev/null @@ -1,5 +0,0 @@ -FROM armaneshaghi/synthseg:latest -#RUN git clone https://github.com/neuropoly/totalsegmentator-mri.git -COPY . /opt/totalsegmentator-mri -WORKDIR /opt/ - diff --git a/README.md b/README.md index 4404233..8be09fb 100644 --- a/README.md +++ b/README.md @@ -1,103 +1,139 @@ -# totalsegmentator-mri -Code for the TotalSegmentator MRI project. +# TotalSegMRI -## Steps to install +Tool for automatic segmentation and labelling of all vertebrae and intervertebral discs (IVDs), spinal cord, and spinal canal. We follow [TotalSegmentator classes](https://github.com/wasserth/TotalSegmentator?tab=readme-ov-file#class-details) with an additional class for IVDs, spinal cord and spinal canal (See list of class [here](#list-of-class)). We used [nnUNet](https://github.com/MIC-DKFZ/nnUNet) as our backbone for model training and inference. -1. Clone this repository - ``` - git clone https://github.com/neuropoly/totalsegmentator-mri.git - ``` +- [Dependencies](#dependencies) +- [Installation](#installation) +- [First Model](#first-model) + - [First Model - Train](#first-model---train) + - [First Model - Inference](#first-model---inference) +- [List of class](#list-of-class) -1. Clone SynthSeg repository - ``` - git clone https://github.com/BBillot/SynthSeg.git - ``` +![Thumbnail](https://github.com/neuropoly/totalsegmentator-mri/assets/36595323/ceca5bb7-f370-477a-8b21-9774853948c6) + +## Dependencies -1. Download [this google folder](https://drive.google.com/drive/folders/11F8q3jhZR0KfHhBpyKygXMo-alTDbp0U?usp=sharing) (the TotalSegmentator example image was downloaded from [here](https://zenodo.org/record/6802614)). +- [Spinal Cord Toolbox (SCT)](https://github.com/neuropoly/spinalcordtoolbox) -1. Create Virtual Environment (Please make sure you're using python 3.8 !!!) +## Installation + +1. Open Terminal in a directory you want to work on. + +1. Create and activate Virtual Environment (Highly recommanded): ``` python -m venv venv + source venv/bin/activate ``` -1. Add SynthSeg to Virtual Environment (If not using bash change '$(pwd)' to the working directory): - ``` - echo "$(pwd)/SynthSeg" > venv/lib/python3.8/site-packages/SynthSeg.pth - ``` +1. Install [PyTorch](https://pytorch.org/get-started/locally/) as described on their website. -1. Activate Virtual Environment +1. Clone and install this repository: ``` - source venv/bin/activate + git clone https://github.com/neuropoly/totalsegmentator-mri.git + python -m pip install -e totalsegmentator-mri ``` 1. Install requirements: ``` - pip install -r SynthSeg/requirements_python3.8.txt + python -m pip install -r totalsegmentator-mri/requirements.txt ``` -## To run scripsts - -`resources/labels.json` - Contain mapping of each mask to unique number. +## First Model +A hybrid approach integrating nnU-Net with an iterative algorithm for segmenting vertebrae, IVDs, spinal cord, and spinal canal. To tackle the challenge of having many classes and class imbalance, we developed a two-step training process. A first model (model 1 - 206) was trained (single input channel: image) to identify 4 classes (IVDs, vertebrae, spinal cord and spinal canal) as well as specific IVDs (C2-C3, C7-T1 and L5-S1) representing key anatomical landmarks along the spine, so 7 classes in total (Figure 1A). The output segmentation was processed using an algorithm that distinguished odd and even IVDs based on the C2-C3, C7-T1 and L5-S1 IVD labels output by the model (Figure 1B). Then, a second nnU-Net model (model 2 - 210) was trained (two input channels: 1=image, 2=odd IVDs), to output 12 classes (Figure 1C). Finally, the output of model 2 was processed in order to assign an individual label value to each vertebrae and IVD in the final segmentation mask (Figure 1D). -`resources/classes.json` - Contain mapping of each mask to class of masks with similar statistics (total 15 classes). +![Figure 1](https://github.com/neuropoly/totalsegmentator-mri/assets/36595323/3958cbc6-a059-4ccf-b3b1-02dbc3a4a62d) -### Option 1 - Run script for all TotalSegmentator labels +**Figure 1**: Illustration of the hybrid method for automatic segmentation of the spine and spinal cord structures. T1w image (A) is used to train model 1, which outputs 7 classes (B). These output labels are processed to extract odd IVDs (C). The T1w and odd IVDs are used as two input channels to train model 2, which outputs 12 classes (D). These output labels are processed to extract individual IVDs and vertebrae (E). -1. Combine all MPRAGE 'blob' masks for each subject into a single segmentation file: - ``` - python totalsegmentator-mri/scripts/combine_masks.py -d TotalSegmentatorMRI_SynthSeg/data/derivatives/manual_masks -o output/ALL_LAB/MP-RAGE_Masks_Combined -m totalsegmentator-mri/resources/labels.json - ``` +### First Model - Train -1. Calculate signal statistics (mean + std) for each masks (group masks into classes of similar statistics): - ``` - python totalsegmentator-mri/scripts/build_intensity_stats.py -d TotalSegmentatorMRI_SynthSeg/data -s output/ALL_LAB/MP-RAGE_Masks_Combined -o output/ALL_LAB/MP-RAGE_priors -m totalsegmentator-mri/resources/labels.json -c totalsegmentator-mri/resources/classes.json - ``` +1. Download the corresponding content from [SPIDER dataset](https://doi.org/10.5281/zenodo.10159290) into 'data/raw/spider/images' and 'data/raw/spider/masks' (you can use `mkdir -p data/raw/spider` to create the folder first). -1. Combine all TotalSegmentator masks for each subject into a single segmentation file: - ``` - python totalsegmentator-mri/scripts/combine_masks.py -d TotalSegmentatorMRI_SynthSeg/Totalsegmentator_dataset -o output/ALL_LAB/TotalSegmentator_Masks_Combined -m totalsegmentator-mri/resources/labels.json --subject-prefix s --subject-subdir segmentations --seg-suffix _ct_seg --output-bids 0 - ``` +1. Make sure `git` and `git-annex` are installed (You can install with `sudo apt-get install git-annex -y`). -1. Create a synthetic image using TotalSegmentator segmentation and the calculated MPRAGE signal statistics: - ``` - python totalsegmentator-mri/scripts/generate_image.py -s output/ALL_LAB/TotalSegmentator_Masks_Combined -p output/ALL_LAB/MP-RAGE_priors -o output/ALL_LAB/MP-RAGE_Synthetic/test1 -n 2 - ``` +1. Extract [data-multi-subject_PAM50_seg.zip](https://drive.google.com/file/d/1Sq38xLHnVxhLr0s1j27ywbeshNUjo3IP) into 'data/bids/data-multi-subject'. -### Option 2 - Run script with TotalSegmentator labels reduced to 15 labels +1. Extract [data-single-subject_PAM50_seg.zip](https://drive.google.com/file/d/1YvuFHL8GDJ5SXlMLORWDjR5SNkDL6TUU) into 'data/bids/data-single-subject'. -To reduce number of labels and group all vertebrae, we use `resources/classes.json` as the main masks mapping when combining masks with combine_masks. This way all masks of the same classes will be mapped to the same label. +1. Extract [whole-spine.zip](https://drive.google.com/file/d/143i0ODmeqohpc4vu5Aa5lnv8LLEyOU0F) (private dataset) into 'data/bids/whole-spine'. -1. Combine all MPRAGE 'blob' masks for each subject into a single segmentation file: +1. Get the required datasets from [Spine Generic Project](https://github.com/spine-generic/): ``` - python totalsegmentator-mri/scripts/combine_masks.py -d TotalSegmentatorMRI_SynthSeg/data/derivatives/manual_masks -o output/15_LAB/MP-RAGE_Masks_Combined -m totalsegmentator-mri/resources/classes.json + source totalsegmentator-mri/run/get_spine_generic_datasets.sh ``` -1. Calculate signal statistics (mean + std) for each masks: +1. Prepares SPIDER datasets in [BIDS](https://bids.neuroimaging.io/) structure: ``` - python totalsegmentator-mri/scripts/build_intensity_stats.py -d TotalSegmentatorMRI_SynthSeg/data -s output/15_LAB/MP-RAGE_Masks_Combined -o output/15_LAB/MP-RAGE_priors -m totalsegmentator-mri/resources/classes.json + source totalsegmentator-mri/run/prepare_spider_bids_datasets.sh ``` -1. Combine all TotalSegmentator masks for each subject into a single segmentation file: +1. Prepares datasets in nnUNetv2 structure: ``` - python totalsegmentator-mri/scripts/combine_masks.py -d TotalSegmentatorMRI_SynthSeg/Totalsegmentator_dataset -o output/15_LAB/TotalSegmentator_Masks_Combined -m totalsegmentator-mri/resources/classes.json --subject-prefix s --subject-subdir segmentations --seg-suffix _ct_seg --output-bids 0 + source totalsegmentator-mri/run/prepare_nnunet_datasets.sh ``` -1. Create a synthetic image using TotalSegmentator segmentation and the calculated MPRAGE signal statistics: +1. Train the model: ``` - python totalsegmentator-mri/scripts/generate_image.py -s output/15_LAB/TotalSegmentator_Masks_Combined -p output/15_LAB/MP-RAGE_priors -o output/15_LAB/MP-RAGE_Synthetic/test1 -n 2 + source totalsegmentator-mri/run/train_nnunet.sh ``` -## Data organization - -As a starting point, a few MPRAGE data are under our private [google folder](https://drive.google.com/drive/folders/1CAkz4ZuxQjWza7GAXhXxTkKcyB9p3yME). -We will follow the BIDS structure: +### First Model - Inference +Run the model on a folder containing the images in .nii.gz format (Make sure to train the model or extract the trained `nnUNet_results` into `data/nnUNet/nnUNet_results` befor running): ``` -├── derivatives -│   └── manual_masks -│   └── sub-errsm37 -│   └── anat -└── sub-errsm37 - └── anat - ├── sub-errsm37_T1w.json - └── sub-errsm37_T1w.nii.gz +source totalsegmentator-mri/run/inference_nnunet.sh INPUT_FOLDER OUTPUT_FOLDER ``` + +## List of class + +|Label|Name| +|:-----|:-----| +| 18 | vertebrae_L5 | +| 19 | vertebrae_L4 | +| 20 | vertebrae_L3 | +| 21 | vertebrae_L2 | +| 22 | vertebrae_L1 | +| 23 | vertebrae_T12 | +| 24 | vertebrae_T11 | +| 25 | vertebrae_T10 | +| 26 | vertebrae_T9 | +| 27 | vertebrae_T8 | +| 28 | vertebrae_T7 | +| 29 | vertebrae_T6 | +| 30 | vertebrae_T5 | +| 31 | vertebrae_T4 | +| 32 | vertebrae_T3 | +| 33 | vertebrae_T2 | +| 34 | vertebrae_T1 | +| 35 | vertebrae_C7 | +| 36 | vertebrae_C6 | +| 37 | vertebrae_C5 | +| 38 | vertebrae_C4 | +| 39 | vertebrae_C3 | +| 40 | vertebrae_C2 | +| 41 | vertebrae_C1 | +| 92 | sacrum | +| 200 | spinal_cord | +| 201 | spinal_canal | +| 202 | disc_L5_S | +| 203 | disc_L4_L5 | +| 204 | disc_L3_L4 | +| 205 | disc_L2_L3 | +| 206 | disc_L1_L2 | +| 207 | disc_T12_L1 | +| 208 | disc_T11_T12 | +| 209 | disc_T10_T11 | +| 210 | disc_T9_T10 | +| 211 | disc_T8_T9 | +| 212 | disc_T7_T8 | +| 213 | disc_T6_T7 | +| 214 | disc_T5_T6 | +| 215 | disc_T4_T5 | +| 216 | disc_T3_T4 | +| 217 | disc_T2_T3 | +| 218 | disc_T1_T2 | +| 219 | disc_C7_T1 | +| 220 | disc_C6_C7 | +| 221 | disc_C5_C6 | +| 222 | disc_C4_C5 | +| 223 | disc_C3_C4 | +| 224 | disc_C2_C3 | diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..c5ce278 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,8 @@ +[build-system] +requires = ["setuptools"] +build-backend = "setuptools.build_meta" + +[project] +name = "totalsegmri" +version = "0.0.1" +dependencies = [] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..be4b9b7 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,10 @@ +tqdm +numpy +nibabel +gryds +torchio +scipy +Pillow +nilearn +SimpleITK +nnunetv2 \ No newline at end of file diff --git a/resources/classes.json b/resources/classes.json deleted file mode 100644 index 5bd1caa..0000000 --- a/resources/classes.json +++ /dev/null @@ -1,107 +0,0 @@ -{ - "background": 0, - "spleen": 1, - "kidney_right": 2, - "kidney_left": 2, - "gallbladder": 3, - "liver": 4, - "stomach": 5, - "aorta": 6, - "inferior_vena_cava": 6, - "portal_vein_and_splenic_vein": 6, - "pancreas": 7, - "adrenal_gland_right": 8, - "adrenal_gland_left": 8, - "lung_upper_lobe_left": 9, - "lung_lower_lobe_left": 9, - "lung_upper_lobe_right": 9, - "lung_middle_lobe_right": 9, - "lung_lower_lobe_right": 9, - "vertebrae_L5": 10, - "vertebrae_L4": 10, - "vertebrae_L3": 10, - "vertebrae_L2": 10, - "vertebrae_L1": 10, - "vertebrae_T12": 10, - "vertebrae_T11": 10, - "vertebrae_T10": 10, - "vertebrae_T9": 10, - "vertebrae_T8": 10, - "vertebrae_T7": 10, - "vertebrae_T6": 10, - "vertebrae_T5": 10, - "vertebrae_T4": 10, - "vertebrae_T3": 10, - "vertebrae_T2": 10, - "vertebrae_T1": 10, - "vertebrae_C7": 10, - "vertebrae_C6": 10, - "vertebrae_C5": 10, - "vertebrae_C4": 10, - "vertebrae_C3": 10, - "vertebrae_C2": 10, - "vertebrae_C1": 10, - "esophagus": 5, - "trachea": 5, - "heart_myocardium": 11, - "heart_atrium_left": 11, - "heart_ventricle_left": 11, - "heart_atrium_right": 11, - "heart_ventricle_right": 11, - "pulmonary_artery": 6, - "brain": 12, - "iliac_artery_left": 6, - "iliac_artery_right": 6, - "iliac_vena_left": 6, - "iliac_vena_right": 6, - "small_bowel": 5, - "duodenum": 5, - "colon": 5, - "rib_left_1": 10, - "rib_left_2": 10, - "rib_left_3": 10, - "rib_left_4": 10, - "rib_left_5": 10, - "rib_left_6": 10, - "rib_left_7": 10, - "rib_left_8": 10, - "rib_left_9": 10, - "rib_left_10": 10, - "rib_left_11": 10, - "rib_left_12": 10, - "rib_right_1": 10, - "rib_right_2": 10, - "rib_right_3": 10, - "rib_right_4": 10, - "rib_right_5": 10, - "rib_right_6": 10, - "rib_right_7": 10, - "rib_right_8": 10, - "rib_right_9": 10, - "rib_right_10": 10, - "rib_right_11": 10, - "rib_right_12": 10, - "humerus_left": 10, - "humerus_right": 10, - "scapula_left": 10, - "scapula_right": 10, - "clavicula_left": 10, - "clavicula_right": 10, - "femur_left": 0, - "femur_right": 0, - "hip_left": 10, - "hip_right": 10, - "sacrum": 10, - "face": 13, - "gluteus_maximus_left": 14, - "gluteus_maximus_right": 14, - "gluteus_medius_left": 14, - "gluteus_medius_right": 14, - "gluteus_minimus_left": 14, - "gluteus_minimus_right": 14, - "autochthon_left": 14, - "autochthon_right": 14, - "iliopsoas_left": 14, - "iliopsoas_right": 14, - "urinary_bladder": 15 -} \ No newline at end of file diff --git a/resources/labels.json b/resources/labels.json deleted file mode 100644 index a7f615e..0000000 --- a/resources/labels.json +++ /dev/null @@ -1,107 +0,0 @@ -{ - "background": 0, - "spleen": 1, - "kidney_right": 2, - "kidney_left": 3, - "gallbladder": 4, - "liver": 5, - "stomach": 6, - "aorta": 7, - "inferior_vena_cava": 8, - "portal_vein_and_splenic_vein": 9, - "pancreas": 10, - "adrenal_gland_right": 11, - "adrenal_gland_left": 12, - "lung_upper_lobe_left": 13, - "lung_lower_lobe_left": 14, - "lung_upper_lobe_right": 15, - "lung_middle_lobe_right": 16, - "lung_lower_lobe_right": 17, - "vertebrae_L5": 18, - "vertebrae_L4": 19, - "vertebrae_L3": 20, - "vertebrae_L2": 21, - "vertebrae_L1": 22, - "vertebrae_T12": 23, - "vertebrae_T11": 24, - "vertebrae_T10": 25, - "vertebrae_T9": 26, - "vertebrae_T8": 27, - "vertebrae_T7": 28, - "vertebrae_T6": 29, - "vertebrae_T5": 30, - "vertebrae_T4": 31, - "vertebrae_T3": 32, - "vertebrae_T2": 33, - "vertebrae_T1": 34, - "vertebrae_C7": 35, - "vertebrae_C6": 36, - "vertebrae_C5": 37, - "vertebrae_C4": 38, - "vertebrae_C3": 39, - "vertebrae_C2": 40, - "vertebrae_C1": 41, - "esophagus": 42, - "trachea": 43, - "heart_myocardium": 44, - "heart_atrium_left": 45, - "heart_ventricle_left": 46, - "heart_atrium_right": 47, - "heart_ventricle_right": 48, - "pulmonary_artery": 49, - "brain": 50, - "iliac_artery_left": 51, - "iliac_artery_right": 52, - "iliac_vena_left": 53, - "iliac_vena_right": 54, - "small_bowel": 55, - "duodenum": 56, - "colon": 57, - "rib_left_1": 58, - "rib_left_2": 59, - "rib_left_3": 60, - "rib_left_4": 61, - "rib_left_5": 62, - "rib_left_6": 63, - "rib_left_7": 64, - "rib_left_8": 65, - "rib_left_9": 66, - "rib_left_10": 67, - "rib_left_11": 68, - "rib_left_12": 69, - "rib_right_1": 70, - "rib_right_2": 71, - "rib_right_3": 72, - "rib_right_4": 73, - "rib_right_5": 74, - "rib_right_6": 75, - "rib_right_7": 76, - "rib_right_8": 77, - "rib_right_9": 78, - "rib_right_10": 79, - "rib_right_11": 80, - "rib_right_12": 81, - "humerus_left": 82, - "humerus_right": 83, - "scapula_left": 84, - "scapula_right": 85, - "clavicula_left": 86, - "clavicula_right": 87, - "femur_left": 0, - "femur_right": 0, - "hip_left": 90, - "hip_right": 91, - "sacrum": 92, - "face": 93, - "gluteus_maximus_left": 94, - "gluteus_maximus_right": 95, - "gluteus_medius_left": 96, - "gluteus_medius_right": 97, - "gluteus_minimus_left": 98, - "gluteus_minimus_right": 99, - "autochthon_left": 100, - "autochthon_right": 101, - "iliopsoas_left": 102, - "iliopsoas_right": 103, - "urinary_bladder": 104 -} \ No newline at end of file diff --git a/run/get_spine_generic_datasets.sh b/run/get_spine_generic_datasets.sh new file mode 100644 index 0000000..0f03f8c --- /dev/null +++ b/run/get_spine_generic_datasets.sh @@ -0,0 +1,52 @@ +#!/bin/bash + +# This script get the data require to train the model from https://github.com/spine-generic/ repository. + +# BASH SETTINGS +# ====================================================================================================================== + +# Uncomment for full verbose +# set -v + +# Immediately exit if error +set -e + +# Exit if user presses CTRL+C (Linux) or CMD+C (OSX) +trap "echo Caught Keyboard Interrupt within script. Exiting now.; exit" INT + +# SCRIPT STARTS HERE +# ====================================================================================================================== + + +# Make sure data/bids exists and enter it +mkdir -p data/bids +cd data/bids + +# Loop over two dataset directories: data-multi-subject and data-single-subject +for ds in data-multi-subject data-single-subject; do + # Clone the dataset from the specified GitHub repository + git clone https://github.com/spine-generic/$ds + cd $ds + + # Remove code and derivatives directories which are not needed for the current task + rm -rf code derivatives + + # Delete all files in the current directory (root of the dataset) except for those named 'sub-*' + find . -mindepth 1 -maxdepth 1 -type f ! -name 'sub-*' -delete + + # Delete all non-anat directories inside each subject's folder + find . -type d -regex '\./sub-[^/]+/.*' ! -regex '\./sub-[^/]+/anat' -exec rm -rf {} \; -prune + + # Within each subject's anat directory, delete all files except T1 and T2 weighted images (both .nii.gz and .json) + find . -type f -regex '\./sub-[^/]+/anat/.*' ! -regex '\./sub-[^/]+/anat/sub-[^_]+_\(T1\|T2\)w\.\(nii\.gz\|json\)' -delete + + # Initialize the current dataset directory as a git-annex repository and download the necessary files + git annex init + git annex get + + # Move back to the parent directory to process the next dataset + cd .. +done + +# Return to the original working directory +cd ../.. diff --git a/run/inference_nnunet.sh b/run/inference_nnunet.sh new file mode 100644 index 0000000..1eceda5 --- /dev/null +++ b/run/inference_nnunet.sh @@ -0,0 +1,83 @@ +#!/bin/bash + +# This script inference tha trained TotalSegMRI nnUNet model. +# this script get the following parameters in the terminal: +# 1'st param: The input folder containing the .nii.gz images to run the model on. +# 2'nd param: The output folder where the models outputs will be stored. + +# BASH SETTINGS +# ====================================================================================================================== + +# Uncomment for full verbose +# set -v + +# Immediately exit if error +set -e + +# Exit if user presses CTRL+C (Linux) or CMD+C (OSX) +trap "echo Caught Keyboard Interrupt within script. Exiting now.; exit" INT + +# SCRIPT STARTS HERE +# ====================================================================================================================== + +INPUT_FOLDER=$1 +OUTPUT_FOLDER=$2 + +# RAM requirement in GB +RAM_REQUIREMENT=8 +# Get the number of CPUs, subtract 1, and ensure the value is at least 1 +JOBS_FOR_CPUS=$(( $(($(nproc) - 1 < 1 ? 1 : $(nproc) - 1 )) )) +# Get the total memory in GB divided by RAM_REQUIREMENT, rounded down to nearest integer, and ensure the value is at least 1 +JOBS_FOR_RAMGB=$(( $(awk -v ram_req="$RAM_REQUIREMENT" '/MemTotal/ {print int($2/1024/1024/ram_req < 1 ? 1 : $2/1024/1024/ram_req)}' /proc/meminfo) )) +# Get the minimum of JOBS_FOR_CPUS and JOBS_FOR_RAMGB +JOBS=$(( JOBS_FOR_CPUS < JOBS_FOR_RAMGB ? JOBS_FOR_CPUS : JOBS_FOR_RAMGB )) + +export nnUNet_def_n_proc=$JOBS +export nnUNet_n_proc_DA=$JOBS + +# Set nnunet params +export nnUNet_raw=data/nnUNet/nnUNet_raw +export nnUNet_preprocessed=data/nnUNet/nnUNet_preprocessed +export nnUNet_results=data/nnUNet/nnUNet_results + +echo "" +echo "Running with the following parameters:" +echo "INPUT_FOLDER=${INPUT_FOLDER}" +echo "OUTPUT_FOLDER=${OUTPUT_FOLDER}" +echo "JOBS=${JOBS}" +echo "" + +# Make output dir with copy of the input +mkdir -p ${OUTPUT_FOLDER}/input +cp ${INPUT_FOLDER}/*.nii.gz ${OUTPUT_FOLDER}/input + +# Add _0000 to inputs if not exists to run nnunet +for f in ${OUTPUT_FOLDER}/input/*.nii.gz; do if [[ $f != *_0000.nii.gz ]]; then mv $f ${f/.nii.gz/_0000.nii.gz}; fi; done + +# Run the first model with postprocessing +nnUNetv2_predict -d 206 -i ${OUTPUT_FOLDER}/input -o ${OUTPUT_FOLDER}/206 -f 0 1 2 3 4 -c 3d_fullres -npp $JOBS -nps $JOBS +nnUNetv2_apply_postprocessing -i ${OUTPUT_FOLDER}/206 -o ${OUTPUT_FOLDER}/206_pp -pp_pkl_file $nnUNet_results/Dataset206_TotalSegMRI/nnUNetTrainer__nnUNetPlans__3d_fullres/crossval_results_folds_0_1_2_3_4/postprocessing.pkl -np $JOBS -plans_json $nnUNet_results/Dataset206_TotalSegMRI/nnUNetTrainer__nnUNetPlans__3d_fullres/plans.json + +# Distinguished odd and even IVDs based on the C2-C3, C7-T1 and L5-S1 IVD labels output by the first model: +# First we will use an iterative algorithm to label IVDs with the definite labels +python totalsegmentator-mri/src/totalsegmri/utils/generate_labels_sequential.py -s ${OUTPUT_FOLDER}/206_pp -o ${OUTPUT_FOLDER}/210_input --output-seg-suffix _0001 --disc-labels 1 2 3 4 --init-disc 2,224 3,219 4,202 --combine-before-label +# Then, we map the IVDs labels to the odd and even IVDs to use as the 2'nd channel of the second model. +python totalsegmentator-mri/src/totalsegmri/utils/map_labels.py -s ${OUTPUT_FOLDER}/210_input -o ${OUTPUT_FOLDER}/210_input -m totalsegmentator-mri/src/totalsegmri/resources/labels_maps/nnunet_210_0001.json --seg-suffix _0001 --output-seg-suffix _0001 + +# For each of the created odd and even IVDs segmentation, copy the original image to use as the 1'st channel intp the second model input folder +for i in ${OUTPUT_FOLDER}/210_input/*; do + cp ${OUTPUT_FOLDER}/input/$(basename ${i//0001.nii.gz/0000.nii.gz}) ${i//0001.nii.gz/0000.nii.gz} +done + +# Run the second model with postprocessing +nnUNetv2_predict -d 210 -i ${OUTPUT_FOLDER}/210_input -o ${OUTPUT_FOLDER}/210 -f 0 1 2 3 4 -c 3d_fullres -npp $JOBS -nps $JOBS +nnUNetv2_apply_postprocessing -i ${OUTPUT_FOLDER}/210 -o ${OUTPUT_FOLDER}/210_pp -pp_pkl_file $nnUNet_results/Dataset210_TotalSegMRI/nnUNetTrainer__nnUNetPlans__3d_fullres/crossval_results_folds_0_1_2_3_4/postprocessing.pkl -np $JOBS -plans_json $nnUNet_results/Dataset210_TotalSegMRI/nnUNetTrainer__nnUNetPlans__3d_fullres/plans.json + +# Use an iterative algorithm to to assign an individual label value to each vertebrae and IVD in the final segmentation mask. +python totalsegmentator-mri/src/totalsegmri/utils/generate_labels_sequential.py -s ${OUTPUT_FOLDER}/210_pp -o ${OUTPUT_FOLDER}/output --disc-labels 2 3 4 5 6 --vertebrea-labels 8 9 10 11 12 --init-disc 4,224 5,219 6,202 --init-vertebrae 10,41 11,34 12,18 + +# Use the spinal cord and canal form the first model output. +python totalsegmentator-mri/src/totalsegmri/utils/map_labels.py -s ${OUTPUT_FOLDER}/206_pp -o ${OUTPUT_FOLDER}/output -m totalsegmentator-mri/src/totalsegmri/resources/labels_maps/nnunet_206_canal.json --add-output + +# Fix csf label to include all non cord spinal canal, this will put the spinal canal label in all the voxels (labeled as a backgroupn) between the spinal canal and the spinal cord. +python totalsegmentator-mri/src/totalsegmri/utils/fix_csf_label.py -s ${OUTPUT_FOLDER}/output -o ${OUTPUT_FOLDER}/output diff --git a/run/prepare_nnunet_datasets.sh b/run/prepare_nnunet_datasets.sh new file mode 100644 index 0000000..41b6375 --- /dev/null +++ b/run/prepare_nnunet_datasets.sh @@ -0,0 +1,80 @@ +#!/bin/bash + +# This script prepares datasets for the TotalSegMRI model in nnUNetv2 structure. + +# BASH SETTINGS +# ====================================================================================================================== + +# Uncomment for full verbose +# set -v + +# Immediately exit if error +set -e + +# Exit if user presses CTRL+C (Linux) or CMD+C (OSX) +trap "echo Caught Keyboard Interrupt within script. Exiting now.; exit" INT + +# SCRIPT STARTS HERE +# ====================================================================================================================== + +for d in 206 210; do + echo "Make nnUNet raw folders ($d)" + mkdir -p data/nnUNet/nnUNet_raw/Dataset${d}_TotalSegMRI/imagesTr + mkdir -p data/nnUNet/nnUNet_raw/Dataset${d}_TotalSegMRI/labelsTr + + echo "Copy dataset file" + cp totalsegmentator-mri/src/totalsegmri/resources/datasets/dataset_${d}.json data/nnUNet/nnUNet_raw/Dataset${d}_TotalSegMRI/dataset.json +done + +# Convert from BIDS to nnUNet dataset, loop over each dataset +for ds in spider data-single-subject data-multi-subject whole-spine; do + echo "Working on $ds" + + # Set segmentation file name based on dataset, 'seg' for 'spider', 'PAM50_seg' otherwise. + [ "$ds" = "spider" ] && seg="seg" || seg="PAM50_seg" + + echo "Copy images and labels into the nnUNet dataset folder" + cp data/bids/$ds/sub-*/anat/sub-*.nii.gz data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/imagesTr + cp data/bids/$ds/derivatives/labels/sub-*/anat/sub-*_$seg.nii.gz data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/labelsTr + + # Format dataset name by removing 'data-' and '-subject' or '-spine' + dsn=${ds/data-/}; dsn=${dsn/-subject/}; dsn=${dsn/-spine/} + + echo "Replace 'sub-' with dataset name" + for f in data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/*Tr/sub-*.nii.gz; do mv $f ${f/sub-/${dsn}_}; done + + echo "Remove _$seg from files name" + for f in data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/labelsTr/*_$seg.nii.gz; do mv $f ${f/_$seg/}; done +done + +echo "Append '_0000' to the images names" +for f in data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/imagesTr/*.nii.gz; do mv $f ${f/.nii.gz/_0000.nii.gz}; done + +echo "Remove images withot segmentation" +for f in data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/imagesTr/*.nii.gz; do if [ ! -f data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/labelsTr/$(basename ${f/_0000.nii.gz/.nii.gz}) ]; then rm $f; fi; done + +echo "Duplicate spider T2Sw X 7, whole X 5 to balance the dataset." +for f in data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/imagesTr/spider*_T2Sw_0000.nii.gz; do for i in {1..6}; do cp $f ${f/_0000.nii.gz/_${i}_0000.nii.gz}; done; done +for f in data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/labelsTr/spider*_T2Sw.nii.gz; do for i in {1..6}; do cp $f ${f/.nii.gz/_${i}.nii.gz}; done; done +for f in data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/imagesTr/whole*_0000.nii.gz; do for i in {1..4}; do cp $f ${f/_0000.nii.gz/_${i}_0000.nii.gz}; done; done +for f in data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/labelsTr/whole*.nii.gz; do for i in {1..4}; do cp $f ${f/.nii.gz/_${i}.nii.gz}; done; done + +echo "Fix csf label to include all non cord spinal canal, this will put the spinal canal label in all the voxels (labeled as a backgroupn) between the spinal canal and the spinal cord." +python totalsegmentator-mri/src/totalsegmri/utils/fix_csf_label.py -s data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/labelsTr -o data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/labelsTr + +echo "Crop images and segmentations in the most anteior voxel of the lowest vertebrae in the image or at the lowest voxel of T12-L1 IVD (for SPIDER dataset)." +for dsn in single multi whole; do + python totalsegmentator-mri/src/totalsegmri/utils/generate_croped_images.py -p ${dsn}_ -i data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/imagesTr -s data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/labelsTr -o data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/imagesTr -g data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/labelsTr +done +python totalsegmentator-mri/src/totalsegmri/utils/generate_croped_images.py -p spider_ -i data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/imagesTr -s data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/labelsTr -o data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/imagesTr -g data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/labelsTr --from-bottom + +echo "Copy from 206 to 210 dataset" +for f in data/nnUNet/nnUNet_raw/Dataset206_TotalSegMRI/*Tr/*.nii.gz; do cp $f ${f/206/210}; done + +echo "Map labels to 210 2'nd channel" +python totalsegmentator-mri/src/totalsegmri/utils/map_labels.py -m totalsegmentator-mri/src/totalsegmri/resources/labels_maps/nnunet_210_0001.json -s data/nnUNet/nnUNet_raw/Dataset210_TotalSegMRI/labelsTr -o data/nnUNet/nnUNet_raw/Dataset210_TotalSegMRI/imagesTr --output-seg-suffix _0001 + +for d in 206 210; do + echo "Map original labels to the dataset specific labels" + python totalsegmentator-mri/src/totalsegmri/utils/map_labels.py -m totalsegmentator-mri/src/totalsegmri/resources/labels_maps/nnunet_${d}.json -s data/nnUNet/nnUNet_raw/Dataset${d}_TotalSegMRI/labelsTr -o data/nnUNet/nnUNet_raw/Dataset${d}_TotalSegMRI/labelsTr +done diff --git a/run/prepare_spider_bids_datasets.sh b/run/prepare_spider_bids_datasets.sh new file mode 100644 index 0000000..af395a8 --- /dev/null +++ b/run/prepare_spider_bids_datasets.sh @@ -0,0 +1,67 @@ +#!/bin/bash + +# This script prepares SPIDER datasets in BIDS structure. + +# BASH SETTINGS +# ====================================================================================================================== + +# Uncomment for full verbose +# set -v + +# Immediately exit if error +set -e + +# Exit if user presses CTRL+C (Linux) or CMD+C (OSX) +trap "echo Caught Keyboard Interrupt within script. Exiting now.; exit" INT + +# SCRIPT STARTS HERE +# ====================================================================================================================== + +# Check if the TotalSegmentator MRI folder exists +if [ ! -d totalsegmentator-mri ]; then + echo "totalsegmentator-mri folder not found in current working directory. Exiting." + echo "Please make sure you clone the repository into the current working directory." + exit 0 +fi + +# Check if the SPIDER dataset images and masks are available +if [ ! -d data/raw/spider/images ] || [ ! -d data/raw/spider/masks ]; then + echo "SPIDER dataset images or masks not found. Please make sure they are available in data/raw/spider." +else + echo "Preparing SPIDER dataset" + + echo "Convert images and masks from .mha to .nii.gz format" + python totalsegmentator-mri/src/totalsegmri/utils/mha2nii.py -i data/raw/spider/images -o data/bids/spider + python totalsegmentator-mri/src/totalsegmri/utils/mha2nii.py -i data/raw/spider/masks -o data/bids/spider/derivatives/labels + + echo "Map SPIDER labels to the labels used in this project" + python totalsegmentator-mri/src/totalsegmri/utils/map_labels.py -m totalsegmentator-mri/src/totalsegmri/resources/labels_maps/spider.json -s data/bids/spider/derivatives/labels -o data/bids/spider/derivatives/labels + + echo "Rename files to follow BIDS naming conventions" + # Add 'sub-' prefix to filenames + for f in data/bids/spider/*.nii.gz; do mv $f ${f/spider\//spider\/sub-}; done + for f in data/bids/spider/derivatives/labels/*.nii.gz; do mv $f ${f/labels\//labels\/sub-}; done + + # Replace 't1' with 'T1w' in filenames + for f in data/bids/spider/*t1.nii.gz; do mv $f ${f/t1/T1w}; done + for f in data/bids/spider/derivatives/labels/*t1.nii.gz; do mv $f ${f/t1/T1w}; done + + # Replace 't2' with 'T2w' in filenames + for f in data/bids/spider/*t2.nii.gz; do mv $f ${f/t2/T2w}; done + for f in data/bids/spider/derivatives/labels/*t2.nii.gz; do mv $f ${f/t2/T2w}; done + + # Replace 't2_SPACE' with 'T2Sw' in filenames + for f in data/bids/spider/*t2_SPACE.nii.gz; do mv $f ${f/t2_SPACE/T2Sw}; done + for f in data/bids/spider/derivatives/labels/*t2_SPACE.nii.gz; do mv $f ${f/t2_SPACE/T2Sw}; done + + # Create anat directories + for f in data/bids/spider/*_*.nii.gz; do mkdir -p ${f/_*.nii.gz/}/anat; done + for f in data/bids/spider/derivatives/labels/*_*.nii.gz; do mkdir -p ${f/_*.nii.gz/}/anat; done + + # Move files into anat directories + for f in data/bids/spider/*_*.nii.gz; do mv $f ${f/_*.nii.gz/}/anat; done + for f in data/bids/spider/derivatives/labels/*_*.nii.gz; do mv $f ${f/_*.nii.gz/}/anat; done + + # Rename segmentation files with '_seg' suffix + for f in data/bids/spider/derivatives/labels/sub-*/anat/sub-*_*.nii.gz; do mv $f ${f/.nii.gz/_seg.nii.gz}; done +fi diff --git a/run/reg2pam50.sh b/run/reg2pam50.sh new file mode 100644 index 0000000..d6d3824 --- /dev/null +++ b/run/reg2pam50.sh @@ -0,0 +1,359 @@ +#!/bin/bash + +# Check for -h option +if [ "$1" = "-h" ]; then + + echo " This script is designed to register spinal cord MRI data to the PAM50 template. " + echo " The data is expected to be organized in the BIDS format. The main steps include:" + echo "" + echo " 1. Create spinal cord segmentation and initial automatic vertebrae label " + echo " (use manual segmentation if available)." + echo " 2. For each subject, use a GUI to manually correct or confirm the vertebrae labels." + echo " 3. Register the image to the PAM50 template and bring the PAM50 segmentation file into the image space." + echo "" + echo " The script takes several command-line arguments for customization:" + echo " -d: BIDS data folder (default: ".")." + echo " -o: Output folder (default: "output")." + echo " -s: PAM50 segmentation file (default: "../PAM50_seg.nii.gz")." + echo " -r: Overwrite existing files (0/1). 0: Do not overwrite, 1: Run again even if files exist (default: 0)." + echo " -c: Cleanup (0/1). 0: Do not Cleanup, 1: Remove all files except the necessary segmentation and label files. (default: 0)." + echo " -l: Log folder (default: "output/logs")." + echo " -j: Number of jobs you want to run in parallel. (default: half the number of cores available on the system)." + echo " -v: Verbose (0/1). 0: Display only errors/warnings, 1: Errors/warnings + info messages (default: 1)." + echo "" + echo " After running the script, a _PAM50_seg.nii.gz file will be created in the SUBJECT output folder " + echo " with the segmentation." + echo "" + echo " The script has three main parts:" + echo "" + echo " 1. Create spinal cord segmentation and initial automatic vertebrae label:" + echo " - Copies the source data to the output folder, preserving the BIDS structure." + echo " - If manual labels are available, they are copied to the output folder." + echo " - Spinal cord segmentation is created using sct_deepseg_sc from the Spinal Cord Toolbox (SCT)." + echo " - Vertebrae labels are created using sct_label_vertebrae." + echo "" + echo " 2. For each subject, use a GUI to manually correct or confirm the vertebrae labels:" + echo " - The script opens SCT's interactive labeling tool (sct_label_utils) for each subject, " + echo " allowing the user to manually correct vertebrae labels." + echo "" + echo " 3. Register the image to the PAM50 template:" + echo " - Registration is performed using sct_register_to_template." + echo " - The PAM50 segmentation file is moved into the image space using sct_apply_transfo." + echo " - Unnecessary intermediate files are removed using the find command, preserving only the final " + echo " segmentation and label files." + echo "" + echo " Note: This script requires the Spinal Cord Toolbox (SCT) to be installed." + echo "" + echo "" + echo " Data organization - follow the BIDS structure:" + echo " data" + echo " ├── derivatives" + echo " │ └── labels" + echo " │ ├── sub-errsm38" + echo " │ │ └── anat" + echo " │ │ ├── sub-errsm38_T1w_labels-disc-manual.json" + echo " │ │ ├── sub-errsm38_T1w_labels-disc-manual.json" + echo " │ │ ├── sub-errsm38_T1w_seg.json" + echo " │ │ └── sub-errsm38_T1w_seg-manual.nii.gz" + echo " │ └── sub-errsm38" + echo " │ └── anat" + echo " │ ├── sub-errsm38_T1w_seg-manual.json" + echo " │ └── sub-errsm38_T1w_seg-manual.nii.gz" + echo " ├── sub-errsm37" + echo " │ └── anat" + echo " │ ├── sub-errsm37_T2w.json" + echo " │ ├── sub-errsm37_T2w.nii.gz" + echo " │ ├── sub-errsm37_T1w.json" + echo " │ └── sub-errsm37_T1w.nii.gz" + echo " └── sub-errsm38" + echo " └── anat" + echo " ├── sub-errsm38_T1w.json" + echo " └── sub-errsm38_T1w.nii.gz" + echo "" + echo " Dependencies" + echo " Spinal Cord Toolbox (SCT) - https://spinalcordtoolbox.com/" + echo "" + echo " Usage:" + echo " ./reg2pam50.sh" + echo " [-d ] \\" + echo " [-o ] \\" + echo " [-s ] \\" + echo " [-r ] \\" + echo " [-c ] \\" + echo " [-l ] \\" + echo " [-j ] \\" + echo " [-v ]" + + exit 0 +fi + +# BASH SETTINGS +# ====================================================================================================================== + +# Uncomment for full verbose +# set -v + +# Immediately exit if error +set -e + +# Exit if user presses CTRL+C (Linux) or CMD+C (OSX) +trap "echo Caught Keyboard Interrupt within script. Exiting now.; exit" INT + +# GET PARAMS +# ====================================================================================================================== + +# Set default values for parameters. +# ---------------------------------------------------------------------------------------------------------------------- +DATA_DIR="." +OUTPUT_DIR="output" +PAM50_SEG="../PAM50_seg.nii.gz" +OVERWRITE=0 +CLEANUP=0 +LOG_DIR="output/logs" +# RAM requirement in GB +RAM_REQUIREMENT=8 +# Get the number of CPUs, subtract 1, and ensure the value is at least 1 +JOBS_FOR_CPUS=$(( $(($(nproc) - 1 < 1 ? 1 : $(nproc) - 1 )) )) +# Get the total memory in GB divided by 10, rounded down to nearest integer, and ensure the value is at least 1 +JOBS_FOR_RAMGB=$(( $(awk -v ram_req="$RAM_REQUIREMENT" '/MemTotal/ {print int($2/1024/1024/ram_req < 1 ? 1 : $2/1024/1024/ram_req)}' /proc/meminfo) )) +# Get the minimum of JOBS_FOR_CPUS and JOBS_FOR_RAMGB +JOBS=$(( JOBS_FOR_CPUS < JOBS_FOR_RAMGB ? JOBS_FOR_CPUS : JOBS_FOR_RAMGB )) +VERBOSE=1 + +# Get command-line parameters to override default values. +# ---------------------------------------------------------------------------------------------------------------------- +while getopts d:o:s:r:c:l:j:v: flag +do + case "${flag}" in + d) DATA_DIR=${OPTARG};; + o) OUTPUT_DIR=${OPTARG};; + s) PAM50_SEG=${OPTARG};; + r) OVERWRITE=${OPTARG};; + c) CLEANUP=${OPTARG};; + l) LOG_DIR=${OPTARG};; + j) JOBS=${OPTARG};; + v) VERBOSE=${OPTARG};; + esac +done + +# Validate the given parameters. +# ---------------------------------------------------------------------------------------------------------------------- + +# Ensure the data directory exists. +if [[ ! -d ${DATA_DIR} ]]; then + echo "Folder not found ${DATA_DIR}" + exit 1 +fi + +# Ensure the output directory exists, creating it if necessary. +mkdir -p ${OUTPUT_DIR} +if [[ ! -d ${OUTPUT_DIR} ]]; then + echo "Folder not found ${OUTPUT_DIR}" + exit 1 +fi + +# Ensure the PAM50 segmentation file exists. +if [[ ! -f ${PAM50_SEG} ]]; then + echo "File not found ${PAM50_SEG}" + exit 1 +fi + +# Ensure the OVERWRITE parameter is either 0 or 1. +if [[ ${OVERWRITE} != 0 ]] && [[ ${OVERWRITE} != 1 ]]; then + echo "Error: -r param must be 0 or 1." + exit 1 +fi + +# Ensure the CLEANUP parameter is either 0 or 1. +if [[ ${CLEANUP} != 0 ]] && [[ ${CLEANUP} != 1 ]]; then + echo "Error: -c param must be 0 or 1." + exit 1 +fi + +# Ensure the log directory exists, creating it if necessary. +mkdir -p ${LOG_DIR} +if [[ ! -d ${LOG_DIR} ]]; then + echo "Folder not found ${LOG_DIR}" + exit 1 +fi + +# Ensure the VERBOSE parameter is either 0 or 1. +if [[ ${VERBOSE} != 0 ]] && [[ ${VERBOSE} != 1 ]]; then + echo "Error: -v param must be 0 or 1." + exit 1 +fi + +# Get full path for all parameters. +# ---------------------------------------------------------------------------------------------------------------------- + +DATA_DIR=$(realpath "${DATA_DIR}") +OUTPUT_DIR=$(realpath "${OUTPUT_DIR}") +PAM50_SEG=$(realpath "${PAM50_SEG}") +LOG_DIR=$(realpath "${LOG_DIR}") + +# Export all parameters. +# ---------------------------------------------------------------------------------------------------------------------- + +export DATA_DIR OUTPUT_DIR PAM50_SEG OVERWRITE CLEANUP LOG_DIR JOBS VERBOSE + +# Print the parameters if VERBOSE is enabled. +# ---------------------------------------------------------------------------------------------------------------------- + +if [[ ${VERBOSE} == 1 ]]; then + echo "" + echo "Running with the following parameters:" + echo "DATA_DIR=${DATA_DIR}" + echo "OUTPUT_DIR=${OUTPUT_DIR}" + echo "PAM50_SEG=${PAM50_SEG}" + echo "OVERWRITE=${OVERWRITE}" + echo "CLEANUP=${CLEANUP}" + echo "LOG_DIR=${LOG_DIR}" + echo "JOBS=${JOBS}" + echo "VERBOSE=${VERBOSE}" + echo "" +fi + +# SCRIPT STARTS HERE +# ====================================================================================================================== + +# Define the contrasts to be processed. +export CONTRASTS="t1 t2" + +# Get running start date and time for log +export start_date=$(date +%Y%m%d%H%M%S) + +# Step 1 - Create spinal cord segmentation and initial automatic vertebrae label. +# ---------------------------------------------------------------------------------------------------------------------- + +# Change to the data directory. +cd ${DATA_DIR} + +seg_sc_and_auto_label_vertebrae () { + + SUBJECT=$1 + + # Copy the source data to the output data folder. + cp -r -u ${DATA_DIR}/${SUBJECT} ${OUTPUT_DIR} + # If there is no anat folder in the subject folder, skip to the next subject. + if [[ ! -d ${OUTPUT_DIR}/${SUBJECT}/anat ]]; then + continue + fi + # Copy manual labels to the output folder, if available. + if [[ -d ${DATA_DIR}/derivatives/labels/${SUBJECT} ]]; then + cp -r -u ${DATA_DIR}/derivatives/labels/${SUBJECT} ${OUTPUT_DIR} + fi + # Change the working directory to the output subject folder. + cd ${OUTPUT_DIR}/${SUBJECT}/anat + # If a manual spinal cord segmentation file exists, use it instead of the automatic one. + for f in sub-*_*w_seg-manual.nii.gz; do + if [[ -e ${f} ]]; then + # Remove '-manual' from segmentation file + mv ${f} "${f%-manual.nii.gz}.nii.gz" || : + fi + done + # Process each contrast. + for c in ${CONTRASTS[@]}; do + if [[ -f ${SUBJECT}_${c^^}w.nii.gz ]]; then + # Get disks label of T1 from T2 and vice versa + cn=$( [ "$c" == "t1" ] && echo "t2" || echo "t1" ) + if [[ ! -f ${SUBJECT}_${c^^}w_labels-disc-manual.nii.gz && -f ${SUBJECT}_${c^^}w_seg.nii.gz && -f ${SUBJECT}_${cn^^}w_labels-disc-manual.nii.gz && -f ${SUBJECT}_${cn^^}w_seg.nii.gz && -f ${SUBJECT}_${cn^^}w.nii.gz ]] ; then + sct_register_multimodal -i ${SUBJECT}_${cn^^}w.nii.gz -d ${SUBJECT}_${c^^}w.nii.gz -iseg ${SUBJECT}_${cn^^}w_seg.nii.gz -dseg ${SUBJECT}_${c^^}w_seg.nii.gz -param step=1,type=seg,algo=slicereg,smooth=1 -v ${VERBOSE} + sct_apply_transfo -i ${SUBJECT}_${cn^^}w_labels-disc-manual.nii.gz -d ${SUBJECT}_${c^^}w.nii.gz -w warp_${SUBJECT}_${cn^^}w2${SUBJECT}_${c^^}w.nii.gz -o ${SUBJECT}_${c^^}w_labels-disc-manual.nii.gz -x label -v ${VERBOSE} + sct_label_vertebrae -i ${SUBJECT}_${c^^}w.nii.gz -s ${SUBJECT}_${c^^}w_seg.nii.gz -discfile ${SUBJECT}_${c^^}w_labels-disc-manual.nii.gz -c ${c} -v ${VERBOSE} -qc ${OUTPUT_DIR}/qc + fi + # Create spinal cord segmentation file if it does not exist or if the OVERWRITE option is enabled. + if [[ ! -f ${SUBJECT}_${c^^}w_seg.nii.gz ]] || [[ ${OVERWRITE} == 1 ]]; then + sct_deepseg_sc -i ${SUBJECT}_${c^^}w.nii.gz -c ${c} -v ${VERBOSE} -qc ${OUTPUT_DIR}/qc + fi + # Label each vertebra if a labeled disc file does not exist and if the manual labels file does not exist, or if the OVERWRITE option is enabled. + if [[ ! -f ${SUBJECT}_${c^^}w_seg_labeled_discs.nii.gz && ! -f ${SUBJECT}_${c^^}w_labels-disc-manual.nii.gz ]] || [[ ${OVERWRITE} == 1 ]]; then + sct_label_vertebrae -i ${SUBJECT}_${c^^}w.nii.gz -s ${SUBJECT}_${c^^}w_seg.nii.gz -c ${c} -v ${VERBOSE} || : + fi + fi + done + +} + +export -f seg_sc_and_auto_label_vertebrae + +# Process all subjects in the data directory in parallel. +ls -d sub-* | parallel -j ${JOBS} "seg_sc_and_auto_label_vertebrae {} >> ${LOG_DIR}/{}_$start_date.log" + +# Step 2 - For each subject, use the GUI to manually correct or confirm the vertebrae labels. +# ---------------------------------------------------------------------------------------------------------------------- + +# Change to the data directory. +cd ${DATA_DIR} + +# Iterate over all subjects in the data directory. +for SUBJECT in sub-*; do + # If there is no anat folder in the subject folder, skip to the next subject. + if [[ ! -d ${OUTPUT_DIR}/${SUBJECT}/anat ]]; then + continue + fi + # Change the working directory to the output subject folder. + cd ${OUTPUT_DIR}/${SUBJECT}/anat + # Process each contrast. + for c in ${CONTRASTS[@]}; do + if [[ -f ${SUBJECT}_${c^^}w.nii.gz ]]; then + # Perform manual correction of vertebrae labels if the manual labels file does not exist or if the OVERWRITE option is enabled. + if [[ ! -f ${SUBJECT}_${c^^}w_labels-disc-manual.nii.gz ]] || [[ ${OVERWRITE} == 1 ]]; then + if [[ -f ${SUBJECT}_${c^^}w_seg_labeled_discs.nii.gz ]]; then + # If the labeled discs file exists, use it as input for the label utils command. + sct_label_utils -i ${SUBJECT}_${c^^}w.nii.gz -create-viewer 1:21 -ilabel ${SUBJECT}_${c^^}w_seg_labeled_discs.nii.gz -v ${VERBOSE} -o ${SUBJECT}_${c^^}w_labels-disc-manual.nii.gz -msg "Place labels at the posterior tip of each inter-vertebral disc. E.g. Label 3: C2/C3, Label 4: C3/C4, etc." || : + else + # If the labeled discs file does not exist, run the label utils command without the input label file. + sct_label_utils -i ${SUBJECT}_${c^^}w.nii.gz -create-viewer 1:21 -v ${VERBOSE} -o ${SUBJECT}_${c^^}w_labels-disc-manual.nii.gz -msg "Place labels at the posterior tip of each inter-vertebral disc. E.g. Label 3: C2/C3, Label 4: C3/C4, etc." || : + fi + fi + fi + done +done + +# Step 3 - Register the image to the PAM50 template. +# ---------------------------------------------------------------------------------------------------------------------- + +# Change to the data directory. +cd ${DATA_DIR} + +register_to_pam50 () { + + SUBJECT=$1 + + # If there is no anat folder in the subject folder, skip to the next subject. + if [[ ! -d ${OUTPUT_DIR}/${SUBJECT}/anat ]]; then + continue + fi + # Change the working directory to the output subject folder. + cd ${OUTPUT_DIR}/${SUBJECT}/anat + # Process each contrast. + for c in ${CONTRASTS[@]}; do + if [[ -f ${SUBJECT}_${c^^}w.nii.gz ]]; then + # Continue if no manual disk labels file is found. + if [[ ! -f ${SUBJECT}_${c^^}w_labels-disc-manual.nii.gz ]]; then + echo "File not found ${SUBJECT}_${c^^}w_labels-disc-manual.nii.gz" + # Register to PAM50 template and create PAM50 segmentation file if it does not exist or if the OVERWRITE option is enabled. + elif [[ ! -f ${SUBJECT}_${c^^}w_PAM50_seg.nii.gz ]] || [[ ${OVERWRITE} == 1 ]]; then + # Register to PAM50 template. + if [[ ! -f ${SUBJECT}_${c^^}w/warp_template2anat.nii.gz ]] || [[ ${OVERWRITE} == 1 ]]; then + # Ensure segmentation in the same space as the image + sct_register_multimodal -i ${SUBJECT}_${c^^}w_seg.nii.gz -d ${SUBJECT}_${c^^}w.nii.gz -identity 1 -x nn + sct_register_to_template -i ${SUBJECT}_${c^^}w.nii.gz -s ${SUBJECT}_${c^^}w_seg_reg.nii.gz -ldisc ${SUBJECT}_${c^^}w_labels-disc-manual.nii.gz -ofolder ${SUBJECT}_${c^^}w -c ${c} -qc ${OUTPUT_DIR}/qc + fi + # Move PAM50 segmentation file into the image space. + sct_apply_transfo -i ${PAM50_SEG} -d ${SUBJECT}_${c^^}w.nii.gz -w ${SUBJECT}_${c^^}w/warp_template2anat.nii.gz -x nn -o ${SUBJECT}_${c^^}w_PAM50_seg.nii.gz + fi + fi + done + # Cleanup - Remove all files except the necessary segmentation and label files. + if [[ ${CLEANUP} == 1 ]]; then + find . -type f -not -regex "\./${SUBJECT}_[^_]+w_\(seg\|labels-disc-manual\|PAM50_seg\)\.nii\.gz" -exec rm -f {} \; + find . -mindepth 1 -maxdepth 1 -type d -exec rm -rf {} \; + fi + +} + +export -f register_to_pam50 + +# Process all subjects in the data directory in parallel. +ls -d sub-* | parallel -j ${JOBS} "register_to_pam50 {} >> ${LOG_DIR}/{}_$start_date.log 2>&1; if [ ! -s ${LOG_DIR}/{}_$start_date.log ]; then rm ${LOG_DIR}/{}_$start_date.log; fi" diff --git a/run/train_nnunet.sh b/run/train_nnunet.sh new file mode 100644 index 0000000..a0b9935 --- /dev/null +++ b/run/train_nnunet.sh @@ -0,0 +1,48 @@ +#!/bin/bash + +# This script train the TotalSegMRI nnUNet models. + +# BASH SETTINGS +# ====================================================================================================================== + +# Uncomment for full verbose +# set -v + +# Immediately exit if error +set -e + +# Exit if user presses CTRL+C (Linux) or CMD+C (OSX) +trap "echo Caught Keyboard Interrupt within script. Exiting now.; exit" INT + +# SCRIPT STARTS HERE +# ====================================================================================================================== + +# RAM requirement in GB +RAM_REQUIREMENT=8 +# Get the number of CPUs, subtract 1, and ensure the value is at least 1 +JOBS_FOR_CPUS=$(( $(($(nproc) - 1 < 1 ? 1 : $(nproc) - 1 )) )) +# Get the total memory in GB divided by RAM_REQUIREMENT, rounded down to nearest integer, and ensure the value is at least 1 +JOBS_FOR_RAMGB=$(( $(awk -v ram_req="$RAM_REQUIREMENT" '/MemTotal/ {print int($2/1024/1024/ram_req < 1 ? 1 : $2/1024/1024/ram_req)}' /proc/meminfo) )) +# Get the minimum of JOBS_FOR_CPUS and JOBS_FOR_RAMGB +JOBS=$(( JOBS_FOR_CPUS < JOBS_FOR_RAMGB ? JOBS_FOR_CPUS : JOBS_FOR_RAMGB )) + +export nnUNet_def_n_proc=$JOBS +export nnUNet_n_proc_DA=$JOBS + +# Set nnunet params +export nnUNet_raw=data/nnUNet/nnUNet_raw +export nnUNet_preprocessed=data/nnUNet/nnUNet_preprocessed +export nnUNet_results=data/nnUNet/nnUNet_results + +for d in 206 210; do + # Preprocess + nnUNetv2_plan_and_preprocess -d $d -c 3d_fullres -npfp $JOBS -np $JOBS --verify_dataset_integrity + + # Train with 5 fold cross-validation + for f in {0..4}; do + nnUNetv2_train $d 3d_fullres $f --npz + done + + # Find best configuration + nnUNetv2_find_best_configuration $d -np $JOBS -c 3d_fullres -f $f +done diff --git a/scripts/build_intensity_stats.py b/scripts/build_intensity_stats.py deleted file mode 100644 index b1bfaa3..0000000 --- a/scripts/build_intensity_stats.py +++ /dev/null @@ -1,262 +0,0 @@ - -import os, sys, argparse, textwrap, json, re -from pathlib import Path -import numpy as np -from SynthSeg.estimate_priors import build_intensity_stats - - -def main(): - - # Description and arguments - parser = argparse.ArgumentParser( - description=textwrap.dedent(f''' - This script uses SynthSeg.estimate_priors.build_intensity_stats to calculate signal statistics (mean + std) ''' - '''for each mask, for all subjects, and saves the rsults to `prior_means.npy` and `prior_stds.npy` files. ''' - '''This script will also generate `labels.npy` - list of unique labels from which signal statistics was ''' - '''calculated. If masks-class-ids was provided `classes.npy` will be generated too, with the class coresponding ''' - '''to the labels in `labels.npy`. The output signal statistics will be calculated for each mask id (from masks-ids ''' - '''argument), with duplicated ids dropped. The output signal statistics will be order by the mask id. - The data is assumed to follow the BIDS structure (or similar): - - data - ├── sub-errsm37 - │ └── anat - │ ├── sub-errsm37_T1w.json - │ └── sub-errsm37_T1w.nii.gz - └── sub-errsm38 - └── anat - ├── sub-errsm38_T1w.json - └── sub-errsm38_T1w.nii.gz - - segmentations - ├── sub-errsm37 - │ └── anat - │ └── sub-errsm37_T1w_seg.nii.gz - └── sub-errsm38 - └── anat - └── sub-errsm38_T1w_seg.nii.gz - - '''), - epilog=textwrap.dedent(''' - Example: - build_intensity_stats -d data -s segmentations -o priors - '''), - formatter_class=argparse.RawTextHelpFormatter - ) - - parser.add_argument( - '--imgs-dir', '-d', type= DirPath(), required=True, - help='The folder where input images are located for each subject.' - ) - parser.add_argument( - '--segs-dir', '-s', type= DirPath(), required=True, - help='The folder where segmentations are located for each subject.' - ) - parser.add_argument( - '--output-dir', '-o', type= DirPath(True), required=True, - help='The folder where output files will be saved to' - ) - parser.add_argument( - '--masks-ids', '-m', type=argparse.FileType('r',encoding='utf-8'), required=True, - help='json file contaning mapping of each mask to a unique number. for example: ' - '{"vertebrae_L3": 2, "organ_spleen": 15}' - ) - parser.add_argument( - '--masks-class-ids', '-c', type=argparse.FileType('r',encoding='utf-8'), default=None, - help='json file contaning mapping of each mask to a class. enables to regroup structures into ' - 'classes of similar intensity statistics. for example: {"vertebrae_L3": 2, "vertebrae_L4": 2}' - ) - parser.add_argument( - '--default-stats', type=int, default=(0, 0, 0, 0), metavar=('mean mean', 'mean std', 'std mean', 'std std'), nargs=4, - help='Default mean (mean+std) and std (mean+std) to use if the calculated value was 0 ' - '(the label is empty). default values is 0.' - ) - parser.add_argument( - '--subject-prefix', type= str, default='sub-', - help='Subject prefix, defaults to "sub-" which is the prefix used for BIDS directories.' - ) - parser.add_argument( - '--subject-subdir', type= str, default='anat', - help='The subfolder inside subject folder contaning the masks, defaults to "anat" which is used for BIDS directories.' - ) - parser.add_argument( - '--img-suffix', type= str, default='_T1w', - help='The suffix for the input images, defaults to "_T1w_seg".' - ) - parser.add_argument( - '--seg-suffix', type= str, default='_T1w_seg', - help='The suffix for the input segmentation, defaults to "_T1w_seg".' - ) - parser.add_argument( - '--verbose', '-v', type= int, default=1, choices=[0, 1], - help='verbose. 0: Display only errors/warnings, 1: Errors/warnings + info messages (default: 1)' - ) - - try: - args = parser.parse_args() - except BaseException as e: - sys.exit() - - # Get args - imgs_path = args.imgs_dir - segs_path = args.segs_dir - output_path = args.output_dir - masks_ids_file = args.masks_ids - masks_class_ids_file = args.masks_class_ids - default_stats = args.default_stats - subject_prefix = args.subject_prefix - subject_subdir = args.subject_subdir - img_suffix = args.img_suffix - seg_suffix = args.seg_suffix - verbose = args.verbose - - if verbose: - print(textwrap.dedent(f''' - Running with the following arguments: - imgs_dir = "{imgs_path}" - segs_dir = "{segs_path}" - output_dir = "{output_path}" - masks_ids = "{masks_ids_file.name if masks_ids_file else ''}" - masks_class_ids = "{masks_class_ids_file.name if masks_class_ids_file else ''}" - default_stats = "{default_stats}" - subject_prefix = "{subject_prefix}" - subject_subdir = "{subject_subdir}" - seg_suffix = "{seg_suffix}" - img_suffix = "{img_suffix}" - verbose = {verbose} - ''')) - - # Get labels from masks-ids json file - estimation_labels = None - if masks_ids_file: - - masks_ids = json.load(masks_ids_file) - masks_ids_file.close() - - # Make estimation_labels unique and sorted as follow: background, non-sided, left, right - background_labels = [masks_ids[k] for k in ['background']] - left_labels = [masks_ids[k] for k in masks_ids if re.match(r'^.*_left(_\d+)?$', k)] - right_labels = [masks_ids[k] for k in masks_ids if re.match(r'^.*_right(_\d+)?$', k)] - non_sided_labels = [v for v in masks_ids.values() if v not in background_labels + left_labels + right_labels] - # Make estimation_labels - estimation_labels = np.array( - background_labels + non_sided_labels + left_labels + right_labels - ) - # Ensure estimation_labels not containing duplicates - estimation_labels = estimation_labels[np.sort(np.unique(estimation_labels, return_index=True)[1])] - - # Get class from masks-class-ids json file - estimation_classes = None - if masks_class_ids_file: - - masks_class_ids = json.load(masks_class_ids_file) - masks_class_ids_file.close() - - # Ensure all masks are in masks_class_ids - if masks_ids is None or estimation_labels is None or not all(m in masks_class_ids for m in masks_ids.keys()): - print(f'Not all masks from {masks_ids_file.name} are in {masks_class_ids.name}') - sys.exit() - - # Create mapping from mask id to mask class id - masks_ids_class_ids_map = {mask_id: masks_class_ids[mask] for mask, mask_id in masks_ids.items()} - - # Create list of masks class ids matching to estimation_labels - estimation_labels_classes = [masks_ids_class_ids_map[mask_id] for mask_id in estimation_labels.tolist()] - - # Generate a map to a new unique and sequential id for each class - seq_estimation_labels_classes_map = {c: i for i, c in enumerate(set(estimation_labels_classes))} - - # Get list of unique and sequential class ids for each label in estimation_labels - estimation_classes = np.array([seq_estimation_labels_classes_map[c] for c in estimation_labels_classes]) - - # Init lists for images and segmantations directories. - segs = [] - imgs = [] - - # Loop over all subjects and look for images and sgmentations files - for sub_path in imgs_path.glob(f'{subject_prefix}*'): - - # Get subject name - subject_name = sub_path.name - if verbose: print(f'Working on {subject_name}') - - # Look for segmentation file and ensure its existance - seg_path = segs_path / subject_name / subject_subdir / f'{subject_name}{seg_suffix}.nii.gz' - - # If segmentation not fount try to find it directly in segs_path (not in subdir) - if not seg_path.exists(): - seg_path = segs_path / f'{subject_name}{seg_suffix}.nii.gz' - - if not seg_path.exists(): - print(f'No segmentation found for {subject_name}') - continue - - # Look for image file and ensure its existance - img_path = imgs_path / subject_name / subject_subdir / f'{subject_name}{img_suffix}.nii.gz' - if not img_path.exists(): - print(f'No image found for {subject_name}') - continue - - # Add image and segentation to the list - segs.append(f'{seg_path}') - imgs.append(f'{img_path}') - - # calculate signal statistics - build_intensity_stats( - list_image_dir=imgs, - list_labels_dir=segs, - estimation_labels=estimation_labels, - estimation_classes=estimation_classes, - result_dir=f'{output_path}', - rescale=True - ) - - # Update default stats from input argument - if default_stats != (0, 0, 0, 0): - for name, vals in zip(('mean', 'std'), np.split(np.array(default_stats), 2)): - fname = f"{output_path / f'prior_{name}s.npy'}" - # Load output file (means or stds) - arr = np.load(fname) - # Update from input argument where the value is 0 - for i, v in enumerate(vals): - arr[i][arr[i]==0] = v - # Save the file - np.save(fname, arr) - - # Save labels and class to file - if estimation_labels is not None: - np.save(output_path / 'labels.npy', estimation_labels) - if estimation_classes is not None: - np.save(output_path / 'classes.npy', estimation_classes) - - -class DirPath(object): - """ - Get path parameter from argparse and return it as pathlib Path object. - - Args: - create (bool): Indicate if the directorie should be created. Default: False. - """ - - def __init__(self, create:bool=False): - self.create = create - - def __call__(self, dir): - - path = Path(dir) - - # Create dir if create was specify - if self.create and not path.exists(): - try: - path.mkdir(parents=True, exist_ok=True) - except: pass - - # Test if path exists - if path.is_dir(): - return path - else: - raise argparse.ArgumentTypeError(f'readable_dir:{path} is not a valid path') - - -if __name__ == '__main__': - main() diff --git a/scripts/combine_masks.py b/scripts/combine_masks.py deleted file mode 100644 index bab5458..0000000 --- a/scripts/combine_masks.py +++ /dev/null @@ -1,218 +0,0 @@ - -import os, sys, argparse, textwrap, re, json -from pathlib import Path -import numpy as np -import nibabel as nib - - -def main(): - - # Description and arguments - parser = argparse.ArgumentParser( - description=textwrap.dedent(f''' - This script combine multiple masks and save the result as a single segmentation file, for each subject. ''' - '''The script sould be provided with a dictionary contained mapping of each mask to its value in the result ''' - '''segmentation. - The data is assumed to follow the BIDS structure (or similar): - - masks - ├── sub-errsm37 - │ └── anat - │ ├── organ_spleen.nii.gz - │ └── vertebrae_L3.nii.gz - └── sub-errsm38 - └── anat - ├── organ_spleen.nii.gz - └── vertebrae_L3.nii.gz - - The output segmentations for this structure will be one of this (depends on the output-bids argument): - - segmentations - ├── sub-errsm37 - │ └── anat - │ └── sub-errsm37_seg.nii.gz - └── sub-errsm38 - └── anat - └── sub-errsm38_seg.nii.gz - - segmentations - ├── sub-errsm37_seg.nii.gz - └── sub-errsm38_seg.nii.gz - - '''), - epilog=textwrap.dedent(''' - Example: - combine_masks -d data -o segmentations -m classes.json - '''), - formatter_class=argparse.RawTextHelpFormatter - ) - - parser.add_argument( - '--masks-dir', '-d', type= DirPath(), required=True, - help='The folder where input masks are located for each subject.' - ) - parser.add_argument( - '--output-dir', '-o', type= DirPath(True), required=True, - help='The folder where output combined segmentations will be saved for each subject.' - ) - parser.add_argument( - '--masks-ids', '-m', type=argparse.FileType('r',encoding='utf-8'), required=True, - help='json file contaning mapping of each mask to a unique number. for example: {"vertebrae_L3": 2, "organ_spleen": 15}' - ) - parser.add_argument( - '--subject-prefix', type= str, default='sub-', - help='Subject prefix, defaults to "sub-" which is the prefix used for BIDS directories.' - ) - parser.add_argument( - '--subject-subdir', type= str, default='anat', - help='The subfolder inside subject folder contaning the masks, defaults to "anat" which is used for BIDS directories.' - ) - parser.add_argument( - '--seg-suffix', type= str, default='_T1w_seg', - help='The suffix for the output segmentation, defaults to "_T1w_seg".' - ) - parser.add_argument( - '--output-bids', type= int, default=1, choices=[0, 1], - help='output-bids. 0: Save output directly in output folder, 1: Save output in BIDS structure (default: 1)' - ) - parser.add_argument( - '--verbose', '-v', type= int, default=1, choices=[0, 1], - help='verbose. 0: Display only errors/warnings, 1: Errors/warnings + info messages (default: 1)' - ) - - try: - args = parser.parse_args() - except BaseException as e: - sys.exit() - - # Get args - masks_path = args.masks_dir - output_path = args.output_dir - masks_ids_file = args.masks_ids - subject_prefix = args.subject_prefix - subject_subdir = args.subject_subdir - seg_suffix = args.seg_suffix - output_bids = args.output_bids - verbose = args.verbose - - if verbose: - print(textwrap.dedent(f''' - Running with the following arguments: - masks_dir = "{masks_path}" - output_dir = "{output_path}" - masks_ids = "{masks_ids_file.name}" - subject_prefix = "{subject_prefix}" - subject_subdir = "{subject_subdir}" - seg_suffix = "{seg_suffix}" - output_bids = {output_bids} - verbose = {verbose} - ''')) - - # Load masks unique values from the json file - masks_ids = json.load(masks_ids_file) - masks_ids_file.close() - - # Loop over all subjects - for sub_path in masks_path.glob(f'{subject_prefix}*'): - - # Ensure its a folder - if not sub_path.is_dir(): - continue - - # Get subject name - subject_name = sub_path.name - - # Combine masks - masks_path = sub_path / subject_subdir - output_file_path = output_path / f'{subject_name}{seg_suffix}.nii.gz' - if output_bids: - # Make output subject name follow the BIDS. - output_subject_name = re.sub(f'^{re.escape(subject_prefix)}', 'sub-', subject_name) - output_file_path = output_path / output_subject_name / 'anat' / f'{output_subject_name}{seg_suffix}.nii.gz' - - if masks_path.exists(): - combine_masks(masks_path, output_file_path, masks_ids) - - -class DirPath(object): - """ - Get path parameter from argparse and return it as pathlib Path object. - - Args: - create (bool): Indicate if the directorie should be created. Default: False. - """ - - def __init__(self, create:bool=False): - self.create = create - - def __call__(self, dir): - - path = Path(dir) - - # Create dir if create was specify - if self.create and not path.exists(): - try: - path.mkdir(parents=True, exist_ok=True) - except: pass - - # Test if path exists - if path.is_dir(): - return path - else: - raise argparse.ArgumentTypeError(f'readable_dir:{path} is not a valid path') - - -def combine_masks(input_masks_dir, output_file, masks_ids, verbose=True): - """ - Combine multiple masks and save the result as a single segmentation file. - The function get dictionary with mapping of each mask to its value in the result segmentation. - - Args: - input_masks_dir (str): Directory where input masks are located. - output_file (str): Output segmentation file path. - masks_ids (dict): Dictionary contain all masks names and their unique id. - The masks should be save to a file with the same name, for example - masks for vertebrae_L3 should be save to vertebrae_L3.nii.gz - verbose (bool, optional): If True, print status updates. Default is True. - """ - - output_data = None - - if verbose: print(f'Looking for masks in {input_masks_dir}') - - input_masks_path = Path(input_masks_dir) - - # Loop over the masks dictionary and look for mask file - for mask_name, id in masks_ids.items(): - - mask_path = input_masks_path / f'{mask_name}.nii.gz' - - # Ensure that the mask exists - if not mask_path.exists(): - continue - - # Load mask - mask = nib.load(mask_path) - mask_data = mask.get_fdata() - - # Convert data type to int to prevent problem with segmentation ids - mask_data = mask_data.astype(np.uint8) - - if output_data is None: - # On the first time init the output segmentation - output_data = mask_data - output_data[np.where(mask_data > 0)] = id - - # Create the result segmentation - combined_seg = nib.Nifti1Image(output_data, mask.affine, mask.header) - combined_seg.set_data_dtype(np.uint8) - - if verbose: print(f'Saving combined segmentation to {output_file}') - - #Make sure output path exists and save results - Path(output_file).parent.mkdir(parents=True, exist_ok=True) - nib.save(combined_seg, output_file) - - -if __name__ == '__main__': - main() diff --git a/scripts/extract_signal_statistics.py b/scripts/extract_signal_statistics.py deleted file mode 100644 index fa17be0..0000000 --- a/scripts/extract_signal_statistics.py +++ /dev/null @@ -1,172 +0,0 @@ - -import os, sys, argparse, textwrap -from pathlib import Path -import numpy as np -import pandas as pd -import nibabel as nib - -def dir_path(path: str) -> Path: - ''' - Get path parameter from argparse and retuen it as pathlib Path object. - ''' - - # Test if path exists - if os.path.isdir(path): - return Path(path) - else: - raise argparse.ArgumentTypeError(f'readable_dir:{path} is not a valid path') - -def main(): - - # Description and arguments - parser = argparse.ArgumentParser( - description=textwrap.dedent(f''' - This script extracts signal statistics (mean + std) for each mask, for each subject, and saves it to a CSV file. The data is assumed to follow the BIDS structure: - - ├── derivatives - │ └── manual_masks - │ ├── sub-errsm37 - │ │ └── anat - │ │ ├── organ_spleen.nii.gz - │ │ └── vertebrae_L3.nii.gz - │ └── sub-errsm38 - │ └── anat - │ ├── organ_spleen.nii.gz - │ └── vertebrae_L3.nii.gz - ├── sub-errsm37 - │ └── anat - │ ├── sub-errsm37_T1w.json - │ └── sub-errsm37_T1w.nii.gz - └── sub-errsm38 - └── anat - ├── sub-errsm38_T1w.json - └── sub-errsm38_T1w.nii.gz - - The output csv will contain row for each subject with mean and std for each mask: - - ---------------------------------------------------------------------------------------- - |subject |organ_spleen_mean |organ_spleen_std |vertebrae_L3_mean |vertebrae_L3_std | - ---------------------------------------------------------------------------------------- - |sub-errsm37 |236.818181818181 |42.1278390252412 |271.25 |37.86010433 | - |sub-errsm38 |252.354838709677 |29.04340288 |223.341234123412 |29.23423432 | - ---------------------------------------------------------------------------------------- - - '''), - epilog=textwrap.dedent(''' - Example: - extract_signal_statistics -d data -o signal_statistics.csv - '''), - formatter_class=argparse.RawTextHelpFormatter - ) - - parser.add_argument( - '--data-dir', '-d', type= dir_path, required=True, - help='The data folder to process.' - ) - parser.add_argument( - '--output-csv', '-o', type=argparse.FileType('w',encoding='utf-8'), required=True, - help='The output csv file.' - ) - parser.add_argument( - '--subject-prefix', type= str, default='sub-', - help='Subject prefix, defaults to "sub-" which is the prefix used for BIDS directories.' - ) - parser.add_argument( - '--image-suffix', type= str, default='_T1w.nii.gz', - help='Image suffix, defaults to "_T1w.nii.gz".' - ) - parser.add_argument( - '--mask-suffix', type= str, default='.nii.gz', - help='Mask suffix Subject prefix, defaults to ".nii.gz".' - ) - parser.add_argument( - '--verbosity', '-v', type= int, default=1, choices=[0, 1], - help='Verbosity. 0: Display only errors/warnings, 1: Errors/warnings + info messages (default: 1)' - ) - - try: - args = parser.parse_args() - except BaseException as e: - sys.exit() - - # Get args - data_dir = args.data_dir - output_csv = args.output_csv - subject_prefix = args.subject_prefix - image_suffix = args.image_suffix - mask_suffix = args.mask_suffix - verbosity = args.verbosity - - if verbosity: - print(textwrap.dedent(f''' - Running with the following arguments: - data_dir = "{data_dir}" - output_csv = "{output_csv.name}" - subject_prefix = "{subject_prefix}" - image_suffix = "{image_suffix}" - mask_suffix = "{mask_suffix}" - verbosity = {verbosity} - ''')) - - # manual_masks directory - manual_masks_dir = data_dir / 'derivatives' / 'manual_masks' - - # Initialize an Array to store results for each subject - subjects_stats = [] - - # Loop across BIDS structure (sub-*) - if verbosity: print(f'Looking for subjects in "{data_dir}"') - for subject_dir in data_dir.glob(f'{subject_prefix}*'): - - if verbosity: print(f'Working on "{subject_dir.name}":') - - # Init dict for subject statistics - subject_stats = {} - subjects_stats.append(subject_stats) - - # Get subject - subject = subject_stats['subject'] = subject_dir.name - - # Get image path - img_path = subject_dir / 'anat' / f'{subject}{image_suffix}' - - # Check if image exist. If not, go to next subject (all values for this subject row will be empty) - if not img_path.exists(): - if verbosity: print(f'Image not found: "{img_path}"') - continue - - if verbosity: print(f'Extracting masks signal statistics for "{img_path.name}"') - - # For each subject (sub-xxx), open the image with nibabel - img = nib.load(img_path) - img_data = img.get_fdata() - - subject_manual_masks_dir = manual_masks_dir / subject / 'anat' - - if verbosity: print(f'Looking for masks in "{subject_manual_masks_dir}"') - - # Loop across the masks under derivatives/manual_masks/sub-xxx/anat/* - for mask_path in sorted((subject_manual_masks_dir).glob(f'*{mask_suffix}')): - - if verbosity: print(f'Extracting signal statistics for "{mask_path.name}"') - - # Get mask name without suffix - mask_name = mask_path.name.replace(mask_suffix, '') - # Open mask - mask = nib.load(mask_path) - mask_data = mask.get_fdata() - - # For each mask, compute the mean and STD and update the dataframe (nan on empty mask) - mask_img_values = img_data[np.where(mask_data)] - subject_stats[f'{mask_name}_mean'] = np.mean(mask_img_values) if mask_img_values.size > 0 else np.nan - subject_stats[f'{mask_name}_std'] = np.std(mask_img_values) if mask_img_values.size > 0 else np.nan - - # Initialize a pandas dataframe (row: subject, column: mask) - df = pd.DataFrame(subjects_stats) - - # Save dataframe as CSV. - if verbosity: print(f'Saving output to "{output_csv.name}"') - df.to_csv(output_csv, index=False, sep=',', lineterminator='\n') - -if __name__ == '__main__': - main() diff --git a/scripts/generate_image.py b/scripts/generate_image.py deleted file mode 100644 index fe96fb9..0000000 --- a/scripts/generate_image.py +++ /dev/null @@ -1,179 +0,0 @@ - -import os, sys, argparse, textwrap, json -from pathlib import Path -import numpy as np -from ext.lab2im import utils -from SynthSeg.brain_generator import BrainGenerator - - -def main(): - - # Description and arguments - parser = argparse.ArgumentParser( - description=textwrap.dedent(f''' - This script uses SynthSeg.brain_generator to generate image from calculated prior signal ''' - '''statistics (mean + std), and saves the rsults `image.nii.gz` and `labels.nii.gz` to output dir. - '''), - epilog=textwrap.dedent(''' - Example: - generate_image -s segmentstions/sub-0287/anat/sub-0287_ct_seg.nii.gz -p priors -o output - '''), - formatter_class=argparse.RawTextHelpFormatter - ) - - parser.add_argument( - '--seg', '-s', type=Path, required=True, - help='The folder where segmentations are located or a segmentation file.' - ) - parser.add_argument( - '--prior-dir', '-p', type=DirPath(), required=True, - help='The folder where calculated prior signal statistics (`prior_means.npy` ,`prior_stds.npy` ' - ' ,`labels.npy` and `classes.npy`) are located.' - ) - parser.add_argument( - '--output-dir', '-o', type=DirPath(True), required=True, - help='The folder where output files (`image.nii.gz` and `labels.nii.gz`) will be saved to' - ) - parser.add_argument( - '--n-images', '-n', type=int, default=1, - help='The number of images to generate (default 1).' - ) - parser.add_argument( - '--verbose', '-v', type=int, default=1, choices=[0, 1], - help='verbose. 0: Display only errors/warnings, 1: Errors/warnings + info messages (default: 1)' - ) - - try: - args = parser.parse_args() - except BaseException as e: - sys.exit() - - # Get args - seg_path = args.seg - prior_path = args.prior_dir - output_path = args.output_dir - n_images = args.n_images - verbose = args.verbose - - if verbose: - print(textwrap.dedent(f''' - Running with the following arguments: - seg = "{seg_path}" - prior_dir = "{prior_path}" - output_dir = "{output_path}" - n_images = "{n_images}" - verbose = {verbose} - ''')) - - # Get priors paths - prior_means_path = prior_path / 'prior_means.npy' - prior_stds_path = prior_path / 'prior_stds.npy' - generation_labels_path = prior_path / 'labels.npy' - generation_classes_path = prior_path / 'classes.npy' - - # Ensure priors exists - if not all([prior_means_path.exists(), prior_stds_path.exists()]): - print(f'`prior_means.npy` or `prior_stds.npy` was not found in {prior_path}') - sys.exit() - - # Get priors values - prior_means = f'{prior_means_path}' - prior_stds = f'{prior_stds_path}' - generation_labels = f'{generation_labels_path}' if generation_labels_path.exists() else None - generation_classes = f'{generation_classes_path}' if generation_classes_path.exists() else None - - if verbose and not generation_labels_path.exists(): - print(f'{generation_labels} was not found') - - if verbose and not generation_classes_path.exists(): - print(f'{generation_classes} was not found') - - # Create synthetic image - brain_generator = BrainGenerator( - labels_dir=f'{seg_path}', - generation_labels=generation_labels, - generation_classes=generation_classes, - prior_means = prior_means, - prior_stds = prior_stds, - n_neutral_labels=None, - output_labels=None, - subjects_prob=None, - batchsize=1, - n_channels=1, - target_res=None, - #target_res=np.array([1., 1., 1.]), - output_shape=200, - output_div_by_n=None, - prior_distributions='uniform', - use_specific_stats_for_channel=False, - mix_prior_and_random=False, - flipping=False, - # scaling_bounds=.2, - scaling_bounds=0.15, - # rotation_bounds=15, - rotation_bounds=10, - # shearing_bounds=.012, - shearing_bounds=0.01, - translation_bounds=False, - nonlin_std=3., - nonlin_scale=.04, - randomise_res=True, - #randomise_res=False, - max_res_iso=4., - max_res_aniso=8., - data_res=None, - #data_res=np.array([1., 1., 1.]), - thickness=None, - #thickness=np.array([1., 1., 1.]), - bias_field_std=.5, - # bias_scale=.025, - bias_scale=0.02, - return_gradients=False - ) - - for n in range(n_images): - - # generate new image and corresponding labels - im, lab = brain_generator.generate_brain() - - # save output image and label map - utils.save_volume( - im, brain_generator.aff, brain_generator.header, - f"{output_path / f'image_{n}.nii.gz'}" - ) - utils.save_volume( - lab, brain_generator.aff, brain_generator.header, - f"{output_path / f'labels_{n}.nii.gz'}" - ) - - -class DirPath(object): - """ - Get path parameter from argparse and return it as pathlib Path object. - - Args: - create (bool): Indicate if the directorie should be created. Default: False. - """ - - def __init__(self, create:bool=False): - self.create = create - - def __call__(self, dir): - - path = Path(dir) - - # Create dir if create was specify - if self.create and not path.exists(): - try: - path.mkdir(parents=True, exist_ok=True) - except: pass - - # Test if path exists - if path.is_dir(): - return path - else: - raise argparse.ArgumentTypeError(f'readable_dir:{path} is not a valid path') - - -if __name__ == '__main__': - main() diff --git a/src/totalsegmri/resources/datasets/dataset_206.json b/src/totalsegmri/resources/datasets/dataset_206.json new file mode 100644 index 0000000..8e65e04 --- /dev/null +++ b/src/totalsegmri/resources/datasets/dataset_206.json @@ -0,0 +1,37 @@ +{ + "channel_names": { + "0": "MRI" + }, + "labels": { + "background": 0, + "disc": [ + 1, + 2, + 3, + 4 + ], + "disc_C2_C3": 2, + "disc_C7_T1": 3, + "disc_L5_S": 4, + "vertebrae": 5, + "canal": [ + 6, + 7, + 8 + ], + "csf": 7, + "cord": 8 + }, + "regions_class_order": [ + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8 + ], + "numTraining": 1869, + "file_ending": ".nii.gz" +} \ No newline at end of file diff --git a/src/totalsegmri/resources/datasets/dataset_210.json b/src/totalsegmri/resources/datasets/dataset_210.json new file mode 100644 index 0000000..c45bd9e --- /dev/null +++ b/src/totalsegmri/resources/datasets/dataset_210.json @@ -0,0 +1,61 @@ +{ + "channel_names": { + "0": "MRI", + "1": "noNorm" + }, + "labels": { + "background": 0, + "disc": [ + 1, + 2, + 3, + 4, + 5, + 6 + ], + "disc_O": 2, + "disc_E": 3, + "disc_C2_C3": 4, + "disc_C7_T1": 5, + "disc_L5_S": 6, + "vertebrae": [ + 7, + 8, + 9, + 10, + 11, + 12 + ], + "vertebrae_O": 8, + "vertebrae_E": 9, + "vertebrae_C1": 10, + "vertebrae_T1": 11, + "vertebrae_L5": 12, + "canal": [ + 13, + 14, + 15 + ], + "csf": 14, + "cord": 15 + }, + "regions_class_order": [ + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15 + ], + "numTraining": 1869, + "file_ending": ".nii.gz" +} \ No newline at end of file diff --git a/src/totalsegmri/resources/labels_maps/mrspineseg.json b/src/totalsegmri/resources/labels_maps/mrspineseg.json new file mode 100644 index 0000000..9219830 --- /dev/null +++ b/src/totalsegmri/resources/labels_maps/mrspineseg.json @@ -0,0 +1,21 @@ +{ + "1": 92, + "2": 18, + "3": 19, + "4": 20, + "5": 21, + "6": 22, + "7": 23, + "8": 24, + "9": 25, + "10": 26, + "11": 202, + "12": 203, + "13": 204, + "14": 205, + "15": 206, + "16": 207, + "17": 208, + "18": 209, + "19": 210 +} \ No newline at end of file diff --git a/src/totalsegmri/resources/labels_maps/nnunet_206.json b/src/totalsegmri/resources/labels_maps/nnunet_206.json new file mode 100644 index 0000000..66ecdd4 --- /dev/null +++ b/src/totalsegmri/resources/labels_maps/nnunet_206.json @@ -0,0 +1,52 @@ +{ + "18": 5, + "19": 5, + "20": 5, + "21": 5, + "22": 5, + "23": 5, + "24": 5, + "25": 5, + "26": 5, + "27": 5, + "28": 5, + "29": 5, + "30": 5, + "31": 5, + "32": 5, + "33": 5, + "34": 5, + "35": 5, + "36": 5, + "37": 5, + "38": 5, + "39": 5, + "40": 5, + "41": 5, + "92": 0, + "200": 8, + "201": 7, + "202": 4, + "203": 1, + "204": 1, + "205": 1, + "206": 1, + "207": 1, + "208": 1, + "209": 1, + "210": 1, + "211": 1, + "212": 1, + "213": 1, + "214": 1, + "215": 1, + "216": 1, + "217": 1, + "218": 1, + "219": 3, + "220": 1, + "221": 1, + "222": 1, + "223": 1, + "224": 2 +} \ No newline at end of file diff --git a/src/totalsegmri/resources/labels_maps/nnunet_206_canal.json b/src/totalsegmri/resources/labels_maps/nnunet_206_canal.json new file mode 100644 index 0000000..70aa598 --- /dev/null +++ b/src/totalsegmri/resources/labels_maps/nnunet_206_canal.json @@ -0,0 +1,4 @@ +{ + "7": 201, + "8": 200 +} diff --git a/src/totalsegmri/resources/labels_maps/nnunet_210.json b/src/totalsegmri/resources/labels_maps/nnunet_210.json new file mode 100644 index 0000000..f037086 --- /dev/null +++ b/src/totalsegmri/resources/labels_maps/nnunet_210.json @@ -0,0 +1,52 @@ +{ + "18": 12, + "19": 8, + "20": 9, + "21": 8, + "22": 9, + "23": 8, + "24": 9, + "25": 8, + "26": 9, + "27": 8, + "28": 9, + "29": 8, + "30": 9, + "31": 8, + "32": 9, + "33": 8, + "34": 11, + "35": 8, + "36": 9, + "37": 8, + "38": 9, + "39": 8, + "40": 9, + "41": 10, + "92": 0, + "200": 15, + "201": 14, + "202": 6, + "203": 3, + "204": 2, + "205": 3, + "206": 2, + "207": 3, + "208": 2, + "209": 3, + "210": 2, + "211": 3, + "212": 2, + "213": 3, + "214": 2, + "215": 3, + "216": 2, + "217": 3, + "218": 2, + "219": 5, + "220": 2, + "221": 3, + "222": 2, + "223": 3, + "224": 4 +} \ No newline at end of file diff --git a/src/totalsegmri/resources/labels_maps/nnunet_210_0001.json b/src/totalsegmri/resources/labels_maps/nnunet_210_0001.json new file mode 100644 index 0000000..2248a7c --- /dev/null +++ b/src/totalsegmri/resources/labels_maps/nnunet_210_0001.json @@ -0,0 +1,52 @@ +{ + "18": 0, + "19": 0, + "20": 0, + "21": 0, + "22": 0, + "23": 0, + "24": 0, + "25": 0, + "26": 0, + "27": 0, + "28": 0, + "29": 0, + "30": 0, + "31": 0, + "32": 0, + "33": 0, + "34": 0, + "35": 0, + "36": 0, + "37": 0, + "38": 0, + "39": 0, + "40": 0, + "41": 0, + "92": 0, + "200": 0, + "201": 0, + "202": 1, + "203": 0, + "204": 1, + "205": 0, + "206": 1, + "207": 0, + "208": 1, + "209": 0, + "210": 1, + "211": 0, + "212": 1, + "213": 0, + "214": 1, + "215": 0, + "216": 1, + "217": 0, + "218": 1, + "219": 0, + "220": 1, + "221": 0, + "222": 1, + "223": 0, + "224": 1 +} diff --git a/src/totalsegmri/resources/labels_maps/spider.json b/src/totalsegmri/resources/labels_maps/spider.json new file mode 100644 index 0000000..60ce5e9 --- /dev/null +++ b/src/totalsegmri/resources/labels_maps/spider.json @@ -0,0 +1,50 @@ +{ + "100": 201, + "201": 202, + "202": 203, + "203": 204, + "204": 205, + "205": 206, + "206": 207, + "207": 208, + "208": 209, + "209": 210, + "210": 211, + "211": 212, + "212": 213, + "213": 214, + "214": 215, + "215": 216, + "216": 217, + "217": 218, + "218": 219, + "219": 220, + "220": 221, + "221": 222, + "222": 223, + "223": 224, + "1": 18, + "2": 19, + "3": 20, + "4": 21, + "5": 22, + "6": 23, + "7": 24, + "8": 25, + "9": 26, + "10": 27, + "11": 28, + "12": 29, + "13": 30, + "14": 31, + "15": 32, + "16": 33, + "17": 34, + "18": 35, + "19": 36, + "20": 37, + "21": 38, + "22": 39, + "23": 40, + "24": 41 +} \ No newline at end of file diff --git a/src/totalsegmri/utils/dirpath.py b/src/totalsegmri/utils/dirpath.py new file mode 100644 index 0000000..8ccfcf9 --- /dev/null +++ b/src/totalsegmri/utils/dirpath.py @@ -0,0 +1,29 @@ +import argparse +from pathlib import Path + +class DirPath: + """ + Get path from argparse and return as Path object. + + Args: + create: Create directory if it doesn't exist + + """ + + def __init__(self, create=False): + self.create = create + + def __call__(self, dir_path): + path = Path(dir_path) + + if self.create and not path.exists(): + try: + path.mkdir(parents=True, exist_ok=True) + except: + pass + + if path.is_dir(): + return path + else: + raise argparse.ArgumentTypeError( + f"readble_dir:{path} is not a valid path") diff --git a/src/totalsegmri/utils/fix_csf_label.py b/src/totalsegmri/utils/fix_csf_label.py new file mode 100644 index 0000000..7be33db --- /dev/null +++ b/src/totalsegmri/utils/fix_csf_label.py @@ -0,0 +1,163 @@ +import sys, argparse, textwrap +import multiprocessing as mp +from tqdm.contrib.concurrent import process_map +from functools import partial +from pathlib import Path +import numpy as np +import nibabel as nib +from totalsegmri.utils.dirpath import DirPath + +def main(): + + # Parse command line arguments + parser = argparse.ArgumentParser( + description=textwrap.dedent(f''' + Fix csf label to include all non cord spinal canal, this will put the spinal canal label in all the voxels (labeled as a backgroupn) between the spinal canal and the spinal cord. + '''), + epilog=textwrap.dedent(''' + Examples: + fix_csf_label -s labels -o labels_fixed + '''), + formatter_class=argparse.RawTextHelpFormatter + ) + + parser.add_argument( + '--segs-dir', '-s', type=DirPath(), required=True, + help='Folder containing input segmentations.' + ) + parser.add_argument( + '--output-dir', '-o', type=DirPath(True), required=True, + help='Folder to save output segmentations.' + ) + parser.add_argument( + '--subject-dir', '-d', type=str, default=None, nargs='?', const='', + help=textwrap.dedent(''' + Is every subject has its oen direcrory. + If this argument will be provided without value it will look for any directory in the segmentation directory. + If value also provided it will be used as a prefix to subject directory (for example "sub-"), defaults to False (no subjet directory). + '''), + ) + parser.add_argument( + '--subject-subdir', '-u', type=str, default='', + help='Subfolder inside subject folder containing masks (for example "anat"), defaults to no subfolder.' + ) + parser.add_argument( + '--prefix', '-p', type=str, default='', + help='File prefix to work on.' + ) + parser.add_argument( + '--seg-suffix', type=str, default='', + help='Segmentation suffix, defaults to "".' + ) + parser.add_argument( + '--output-seg-suffix', type=str, default='', + help='Suffix for output segmentation, defaults to "".' + ) + parser.add_argument( + '--cord-label', type=int, default=200, + help='Label used for spinal cord, defaults to 200.' + ) + parser.add_argument( + '--csf-label', type=int, default=201, + help='Label used for csf, defaults to 201.' + ) + parser.add_argument( + '--max-workers', '-w', type=int, default=min(32, mp.cpu_count() + 4), + help='Max worker to run in parallel proccess, defaults to min(32, mp.cpu_count() + 4).' + ) + parser.add_argument( + '--verbose', '-v', type=int, default=1, choices=[0, 1], + help='Verbosity level. 0: Errors/warnings only, 1: Errors/warnings + info (default: 1)' + ) + + try: + args = parser.parse_args() + except BaseException as e: + sys.exit() + + # Get arguments + seg_path = args.seg_dir + output_path = args.output_dir + subject_dir = args.subject_dir + subject_subdir = args.subject_subdir + prefix = args.prefix + seg_suffix = args.seg_suffix + output_seg_suffix = args.output_seg_suffix + cord_label = args.cord_label + csf_label = args.csf_label + max_workers = args.max_workers + verbose = args.verbose + + if verbose: + print(textwrap.dedent(f''' + Running {Path(__file__).stem} with the following params: + seg_dir = "{seg_path}" + output_dir = "{output_path}" + subject_dir = "{subject_dir}" + subject_subdir = "{subject_subdir}" + prefix = "{prefix}" + seg_suffix = "{seg_suffix}" + output_seg_suffix = "{output_seg_suffix}" + cord_label = "{cord_label}" + csf_label = "{csf_label}" + max_workers = "{max_workers}" + verbose = {verbose} + ''')) + + glob_pattern = "" + if subject_dir is not None: + glob_pattern += f"{subject_dir}*/" + if len(subject_subdir) > 0: + glob_pattern += f"{subject_subdir}/" + glob_pattern += f'{prefix}*{seg_suffix}.nii.gz' + + # Process the NIfTI image and segmentation files + segs_path_list = list(seg_path.glob(glob_pattern)) + + # Create a partially-applied function with the extra arguments + partial_fix_csf_label = partial(fix_csf_label, cord_label=cord_label, csf_label=csf_label, output_path=output_path, seg_suffix=seg_suffix, output_seg_suffix=output_seg_suffix) + + with mp.Pool() as pool: + process_map(partial_fix_csf_label, segs_path_list, max_workers=max_workers) + + +def fix_csf_label(seg_path, cord_label, csf_label, output_path, seg_suffix, output_seg_suffix): + + output_seg_path = output_path / seg_path.name.replace(f'{seg_suffix}.nii.gz', f'{output_seg_suffix}.nii.gz') + + # Load segmentation + seg = nib.load(seg_path) + seg_data = seg.get_fdata() + + # Convert data to uint8 to avoid issues with segmentation IDs + seg_data = seg_data.astype(np.uint8) + + # Create an array of x indices + x_indices = np.broadcast_to(np.arange(seg_data.shape[0])[..., np.newaxis, np.newaxis], seg_data.shape) + # Create an array of y indices + y_indices = np.broadcast_to(np.arange(seg_data.shape[1])[..., np.newaxis], seg_data.shape) + + canal_mask = np.isin(seg_data, [cord_label, csf_label]) + canal_mask_min_x = np.min(np.where(canal_mask, x_indices, np.inf), axis=0)[np.newaxis, ...] + canal_mask_max_x = np.max(np.where(canal_mask, x_indices, -np.inf), axis=0)[np.newaxis, ...] + canal_mask_min_y = np.min(np.where(canal_mask, y_indices, np.inf), axis=1)[:, np.newaxis, :] + canal_mask_max_y = np.max(np.where(canal_mask, y_indices, -np.inf), axis=1)[:, np.newaxis, :] + canal_mask = \ + (canal_mask_min_x <= x_indices) & \ + (x_indices <= canal_mask_max_x) & \ + (canal_mask_min_y <= y_indices) & \ + (y_indices <= canal_mask_max_y) + seg_data[canal_mask & (seg_data != cord_label)] = csf_label + + # Create result segmentation + mapped_seg = nib.Nifti1Image(seg_data, seg.affine, seg.header) + mapped_seg.set_data_dtype(np.uint8) + + # Make sure output directory exists + output_seg_path.parent.mkdir(parents=True, exist_ok=True) + + # Save mapped segmentation + nib.save(mapped_seg, output_seg_path) + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/totalsegmri/utils/generate_augmentations.py b/src/totalsegmri/utils/generate_augmentations.py new file mode 100644 index 0000000..7a9bced --- /dev/null +++ b/src/totalsegmri/utils/generate_augmentations.py @@ -0,0 +1,386 @@ +import sys, argparse, textwrap +import multiprocessing as mp +from functools import partial +from tqdm.contrib.concurrent import process_map +from pathlib import Path +import nibabel as nib +import numpy as np +import torchio as tio +import gryds +from totalsegmri.utils.dirpath import DirPath + + +def main(): + + # Description and arguments + parser = argparse.ArgumentParser( + description=textwrap.dedent(f''' + This script processes NIfTI (Neuroimaging Informatics Technology Initiative) image and segmentation files. + It apply transformation on the image and the segmentation to make augmented image.''' + ), + formatter_class=argparse.RawTextHelpFormatter + ) + parser.add_argument( + '--images-dir', '-i', type=DirPath(), required=True, + help='The folder where input NIfTI images files are located (required).' + ) + parser.add_argument( + '--segs-dir', '-s', type=DirPath(), required=True, + help='The folder where input NIfTI segmentation files are located (required).' + ) + parser.add_argument( + '--output-images-dir', '-o', type=DirPath(True), required=True, + help='The folder where output combined JPG images will be saved (required).' + ) + parser.add_argument( + '--output-segs-dir', '-g', type=DirPath(True), required=True, + help='The folder where output combined JPG images will be saved (required).' + ) + parser.add_argument( + '--subject-dir', '-d', type=str, default=None, nargs='?', const='', + help=textwrap.dedent(''' + Is every subject has its oen direcrory. + If this argument will be provided without value it will look for any directory in the segmentation directory. + If value also provided it will be used as a prefix to subject directory, defaults to False (no subjet directory). + '''), + ) + parser.add_argument( + '--subject-subdir', '-u', type=str, default='', + help='Subfolder inside subject folder containing masks, defaults to no subfolder.' + ) + parser.add_argument( + '--prefix', '-p', type=str, default='', + help='File prefix to work on.' + ) + parser.add_argument( + '--image-suffix', type=str, default='_0000', + help='Image suffix, defaults to "_0000".' + ) + parser.add_argument( + '--seg-suffix', type=str, default='', + help='Segmentation suffix, defaults to "".' + ) + parser.add_argument( + '--output-image-suffix', type=str, default='_0000', + help='Image suffix for output, defaults to "_0000".' + ) + parser.add_argument( + '--output-seg-suffix', type=str, default='', + help='Segmentation suffix for output, defaults to "".' + ) + parser.add_argument( + '--augmentations-per-image', '-n', type=int, default=7, + help='Number of augmentation images to generate. Default is 7.' + ) + parser.add_argument( + '--max-workers', '-w', type=int, default=min(32, mp.cpu_count() + 4), + help='Max worker to run in parallel proccess, defaults to min(32, mp.cpu_count() + 4).' + ) + parser.add_argument( + '--verbose', '-v', type=int, default=1, choices=[0, 1], + help='verbose. 0: Display only errors/warnings, 1: Errors/warnings + info messages. Default is 1.' + ) + + # Parse the command-line arguments + try: + args = parser.parse_args() + except BaseException as e: + sys.exit() + + # Get the command-line argument values + images_path = args.images_dir + segs_path = args.segs_dir + output_images_path = args.output_images_dir + output_segs_path = args.output_segs_dir + subject_dir = args.subject_dir + subject_subdir = args.subject_subdir + prefix = args.prefix + image_suffix = args.image_suffix + seg_suffix = args.seg_suffix + output_image_suffix = args.output_image_suffix + output_seg_suffix = args.output_seg_suffix + augmentations_per_image = args.augmentations_per_image + max_workers = args.max_workers + verbose = args.verbose + + # Print the argument values if verbose is enabled + if verbose: + print(textwrap.dedent(f''' + Running {Path(__file__).stem} with the following params: + images_path = "{images_path}" + segs_path = "{segs_path}" + output_images_path = "{output_images_path}" + output_segs_path = "{output_segs_path}" + subject_dir = "{subject_dir}" + subject_subdir = "{subject_subdir}" + prefix = "{prefix}" + image_suffix = "{image_suffix}" + seg_suffix = "{seg_suffix}" + output_image_suffix = "{output_image_suffix}" + output_seg_suffix = "{output_seg_suffix}" + augmentations_per_image = "{augmentations_per_image}" + max_workers = "{max_workers}" + verbose = "{verbose}" + ''')) + + glob_pattern = "" + if subject_dir is not None: + glob_pattern += f"{subject_dir}*/" + if len(subject_subdir) > 0: + glob_pattern += f"{subject_subdir}/" + glob_pattern += f'{prefix}*{image_suffix}.nii.gz' + + # Process the NIfTI image and segmentation files + images_path_list = list(images_path.glob(glob_pattern)) + + # Create a partially-applied function with the extra arguments + partial_generate_augmentations = partial( + generate_augmentations, + images_path=images_path, + segs_path=segs_path, + output_images_path=output_images_path, + output_segs_path=output_segs_path, + image_suffix=image_suffix, + seg_suffix=seg_suffix, + output_image_suffix=output_image_suffix, + output_seg_suffix=output_seg_suffix, + augmentations_per_image=augmentations_per_image, + ) + + with mp.Pool() as pool: + process_map(partial_generate_augmentations, images_path_list, max_workers=max_workers) + + +def generate_augmentations( + image_path, + images_path, + segs_path, + output_images_path, + output_segs_path, + image_suffix, + seg_suffix, + output_image_suffix, + output_seg_suffix, + augmentations_per_image, + ): + + seg_path = segs_path / image_path.relative_to(images_path).parent / image_path.name.replace(f'{image_suffix}.nii.gz', f'{seg_suffix}.nii.gz') + + if not seg_path.is_file(): + print(f'Segmentation file not found: {seg_path}') + return + + image = nib.load(image_path) + image_data = image.get_fdata() + seg = nib.load(seg_path) + seg_data = seg.get_fdata().astype(np.uint8) + + for i in range(augmentations_per_image): + output_image_path = output_images_path / image_path.relative_to(images_path).parent / image_path.name.replace(f'{image_suffix}.nii.gz', f'_a{i}{output_image_suffix}.nii.gz') + output_seg_path = output_segs_path / image_path.relative_to(images_path).parent / seg_path.name.replace(f'{seg_suffix}.nii.gz', f'_a{i}{output_seg_suffix}.nii.gz') + + augs = [] + + # Randomly add aug_labels2image, aug_flip, aug_inverse with respective probabilities + if np.random.rand() < 0.1: + augs.append(aug_labels2image) + if np.random.rand() < 0.3: + augs.append(aug_flip) + if np.random.rand() < 0.3: + augs.append(aug_inverse) + + augs.extend(np.random.choice([ + np.random.choice([ + aug_bspline, + aug_aff, + aug_elastic, + ]), + np.random.choice([ + aug_log, + aug_sqrt, + aug_exp, + aug_sin, + aug_sig, + aug_gamma, + ]), + np.random.choice([ + aug_motion, + aug_ghosting, + aug_spike, + aug_bias_field, + aug_blur, + aug_noise, + ]), + ], + np.random.choice(range(1,4)), + False + )) + + if np.random.rand() < 0.7: + augs.append(aug_anisotropy) + + # Augment the images + output_image_data, output_seg_data = image_data, seg_data + for a in augs: + output_image_data, output_seg_data = a(output_image_data, output_seg_data) + # output_seg_data = output_seg_data.astype(np.uint8) + + # Create result + output_image = nib.Nifti1Image(output_image_data, image.affine, image.header) + output_seg = nib.Nifti1Image(output_seg_data, seg.affine, seg.header) + output_seg.set_data_dtype(np.uint8) + + # Make sure output directory exists + output_image_path.parent.mkdir(parents=True, exist_ok=True) + output_seg_path.parent.mkdir(parents=True, exist_ok=True) + + # Save mapped segmentation + nib.save(output_image, output_image_path) + nib.save(output_seg, output_seg_path) + +def aug_transform(img, seg, transform): + # Compute original mean, std and min/max values + original_mean, original_std = np.mean(img), np.std(img) + img_min, img_max = np.min(img), np.max(img) + + # Generate random mean and std + random_mean = np.random.uniform(img_min, img_max) + random_std = np.random.uniform(0.5 * original_std, 2 * original_std) + + # Normlize + img = (img - original_mean) / original_std + img = (img - img.min()) / (img.max() - img.min()) + + # Transform + img = transform(img) + + # Redistribution + img = img * random_std + random_mean + + # Return to original range + img = np.interp(img, (img_min, img_max), (img_min, img_max)) + + return img, seg + +def aug_log(img, seg): + return aug_transform(img, seg, np.log) + +def aug_sqrt(img, seg): + return aug_transform(img, seg, np.sqrt) + +def aug_sin(img, seg): + return aug_transform(img, seg, np.sin) + +def aug_exp(img, seg): + return aug_transform(img, seg, np.exp) + +def aug_sig(img, seg): + return aug_transform(img, seg, lambda x:1/(1 + np.exp(-x))) + +def aug_inverse(img, seg): + img = img.min() + img.max() - img + return img, seg + +def aug_bspline(img, seg): + grid = np.random.rand(3, 3, 3, 3) + + bspline = gryds.BSplineTransformation((grid-.5)/5) + grid[:,0] += ((grid[:,0] > 0) * 2 - 1) * .9 # Increase the effect on the Y-axis + return gryds.Interpolator(img).transform(bspline), gryds.Interpolator(seg, order=0).transform(bspline).astype(np.uint8) + +def aug_flip(img, seg): + subject = tio.RandomFlip(axes=('LR',))(tio.Subject( + image=tio.ScalarImage(tensor=np.expand_dims(img, axis=0)), + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + +def aug_aff(img, seg): + subject = tio.RandomAffine()(tio.Subject( + image=tio.ScalarImage(tensor=np.expand_dims(img, axis=0)), + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + +def aug_elastic(img, seg): + subject = tio.RandomElasticDeformation(max_displacement=40)(tio.Subject( + image=tio.ScalarImage(tensor=np.expand_dims(img, axis=0)), + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + +def aug_anisotropy(img, seg): + subject = tio.RandomAnisotropy()(tio.Subject( + image=tio.ScalarImage(tensor=np.expand_dims(img, axis=0)), + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + +def aug_motion(img, seg): + subject = tio.RandomMotion()(tio.Subject( + image=tio.ScalarImage(tensor=np.expand_dims(img, axis=0)), + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + +def aug_ghosting(img, seg): + subject = tio.RandomGhosting()(tio.Subject( + image=tio.ScalarImage(tensor=np.expand_dims(img, axis=0)), + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + +def aug_spike(img, seg): + subject = tio.RandomSpike()(tio.Subject( + image=tio.ScalarImage(tensor=np.expand_dims(img, axis=0)), + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + +def aug_bias_field(img, seg): + subject = tio.RandomBiasField()(tio.Subject( + image=tio.ScalarImage(tensor=np.expand_dims(img, axis=0)), + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + +def aug_blur(img, seg): + subject = tio.RandomBlur()(tio.Subject( + image=tio.ScalarImage(tensor=np.expand_dims(img, axis=0)), + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + +def aug_noise(img, seg): + original_mean, original_std = np.mean(img), np.std(img) + img = (img - original_mean) / original_std + subject = tio.RandomNoise()(tio.Subject( + image=tio.ScalarImage(tensor=np.expand_dims(img, axis=0)), + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + img = img * original_std + original_mean + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + +def aug_swap(img, seg): + subject = tio.RandomSwap()(tio.Subject( + image=tio.ScalarImage(tensor=np.expand_dims(img, axis=0)), + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + +def aug_labels2image(img, seg): + subject = tio.RandomLabelsToImage(label_key="seg", image_key="image")(tio.Subject( + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + +def aug_gamma(img, seg): + subject = tio.RandomGamma((-1, 1))(tio.Subject( + image=tio.ScalarImage(tensor=np.expand_dims(img, axis=0)), + seg=tio.LabelMap(tensor=np.expand_dims(seg, axis=0)) + )) + return subject.image.data.squeeze().numpy(), subject.seg.data.squeeze().numpy().astype(np.uint8) + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/totalsegmri/utils/generate_croped_images.py b/src/totalsegmri/utils/generate_croped_images.py new file mode 100644 index 0000000..f48257a --- /dev/null +++ b/src/totalsegmri/utils/generate_croped_images.py @@ -0,0 +1,224 @@ +import sys, argparse, textwrap, tempfile, shutil, subprocess +from pathlib import Path +import multiprocessing as mp +from functools import partial +from pathlib import Path +from tqdm.contrib.concurrent import process_map +import nibabel as nib +import numpy as np +from totalsegmri.utils.dirpath import DirPath + + +def main(): + + # Description and arguments + parser = argparse.ArgumentParser( + description=textwrap.dedent(f''' + This script processes NIfTI (Neuroimaging Informatics Technology Initiative) image and segmentation files. + It crop images and segmentations in the most anteior voxel the lowest vertebrae in the image or at the lowest voxel of T12-L1 IVD.''' + ), + formatter_class=argparse.RawTextHelpFormatter + ) + parser.add_argument( + '--images-dir', '-i', type=DirPath(), required=True, + help='The folder where input NIfTI images files are located (required).' + ) + parser.add_argument( + '--segs-dir', '-s', type=DirPath(), required=True, + help='The folder where input NIfTI segmentation files are located (required).' + ) + parser.add_argument( + '--output-images-dir', '-o', type=DirPath(True), required=True, + help='The folder where output combined JPG images will be saved (required).' + ) + parser.add_argument( + '--output-segs-dir', '-g', type=DirPath(True), required=True, + help='The folder where output combined JPG images will be saved (required).' + ) + parser.add_argument( + '--subject-dir', '-d', type=str, default=None, nargs='?', const='', + help=textwrap.dedent(''' + Is every subject has its oen direcrory. + If this argument will be provided without value it will look for any directory in the segmentation directory. + If value also provided it will be used as a prefix to subject directory, defaults to False (no subjet directory). + '''), + ) + parser.add_argument( + '--subject-subdir', '-u', type=str, default='', + help='Subfolder inside subject folder containing masks, defaults to no subfolder.' + ) + parser.add_argument( + '--prefix', '-p', type=str, default='', + help='File prefix to work on.' + ) + parser.add_argument( + '--image-suffix', type=str, default='_0000', + help='Image suffix, defaults to "_0000".' + ) + parser.add_argument( + '--seg-suffix', type=str, default='', + help='Segmentation suffix, defaults to "".' + ) + parser.add_argument( + '--output-image-suffix', type=str, default='_0000', + help='Image suffix for output, defaults to "_0000".' + ) + parser.add_argument( + '--output-seg-suffix', type=str, default='', + help='Segmentation suffix for output, defaults to "".' + ) + parser.add_argument( + '--from-bottom', action="store_true", default=False, + help='Crop at the lowest voxel of T12-L1 IVD. Default: Crop in the most anteior voxel the lowest vertebrae in the image' + ) + parser.add_argument( + '--max-workers', '-w', type=int, default=min(32, mp.cpu_count() + 4), + help='Max worker to run in parallel proccess, defaults to min(32, mp.cpu_count() + 4).' + ) + parser.add_argument( + '--verbose', '-v', type=int, default=1, choices=[0, 1], + help='verbose. 0: Display only errors/warnings, 1: Errors/warnings + info messages. Default is 1.' + ) + + # Parse the command-line arguments + try: + args = parser.parse_args() + except BaseException as e: + sys.exit() + + # Get the command-line argument values + images_path = args.images_dir + segs_path = args.segs_dir + output_images_path = args.output_images_dir + output_segs_path = args.output_segs_dir + subject_dir = args.subject_dir + subject_subdir = args.subject_subdir + prefix = args.prefix + image_suffix = args.image_suffix + seg_suffix = args.seg_suffix + output_image_suffix = args.output_image_suffix + output_seg_suffix = args.output_seg_suffix + from_bottom = args.from_bottom + max_workers = args.max_workers + verbose = args.verbose + + # Print the argument values if verbose is enabled + if verbose: + print(textwrap.dedent(f''' + Running {Path(__file__).stem} with the following params: + images_path = "{images_path}" + segs_path = "{segs_path}" + output_images_path = "{output_images_path}" + output_segs_path = "{output_segs_path}" + subject_dir = "{subject_dir}" + subject_subdir = "{subject_subdir}" + prefix = "{prefix}" + image_suffix = "{image_suffix}" + seg_suffix = "{seg_suffix}" + output_image_suffix = "{output_image_suffix}" + output_seg_suffix = "{output_seg_suffix}" + from_bottom = "{from_bottom}" + max_workers = "{max_workers}" + verbose = "{verbose}" + ''')) + + glob_pattern = "" + if subject_dir is not None: + glob_pattern += f"{subject_dir}*/" + if len(subject_subdir) > 0: + glob_pattern += f"{subject_subdir}/" + glob_pattern += f'{prefix}*{image_suffix}.nii.gz' + + # Process the NIfTI image and segmentation files + images_path_list = list(images_path.glob(glob_pattern)) + + # Create a partially-applied function with the extra arguments + partial_generate_croped_images = partial( + generate_croped_images, + images_path=images_path, + segs_path=segs_path, + output_images_path=output_images_path, + output_segs_path=output_segs_path, + image_suffix=image_suffix, + seg_suffix=seg_suffix, + output_image_suffix=output_image_suffix, + output_seg_suffix=output_seg_suffix, + from_bottom=from_bottom, + ) + + with mp.Pool() as pool: + process_map(partial_generate_croped_images, images_path_list, max_workers=max_workers) + + +def generate_croped_images( + image_path, + images_path, + segs_path, + output_images_path, + output_segs_path, + image_suffix, + seg_suffix, + output_image_suffix, + output_seg_suffix, + from_bottom, + ): + + seg_path = segs_path / image_path.relative_to(images_path).parent / image_path.name.replace(f'{image_suffix}.nii.gz', f'{seg_suffix}.nii.gz') + + if not seg_path.is_file(): + print(f'Segmentation file not found: {seg_path}') + return + + output_image_path = output_images_path / image_path.relative_to(images_path).parent / image_path.name.replace(f'{image_suffix}.nii.gz', f'{output_image_suffix}.nii.gz') + output_seg_path = output_segs_path / image_path.relative_to(images_path).parent / seg_path.name.replace(f'{seg_suffix}.nii.gz', f'{output_seg_suffix}.nii.gz') + + temp_path = Path(tempfile.mkdtemp()) + output_image_path.parent.mkdir(parents=True, exist_ok=True) + output_seg_path.parent.mkdir(parents=True, exist_ok=True) + + try: + + # Copy files to tmp in canonial orientation + run_command(f'sct_image -i {image_path} -setorient LPI -o {temp_path}/img.nii.gz') + run_command(f'sct_image -i {seg_path} -setorient LPI -o {temp_path}/seg.nii.gz') + + # Ensure img and seg in the same space + run_command(f'sct_register_multimodal -i {temp_path}/seg.nii.gz -d {temp_path}/img.nii.gz -identity 1 -x nn -o {temp_path}/seg.nii.gz') + + seg = nib.load(f'{temp_path}/seg.nii.gz') + seg_data = seg.get_fdata().astype(np.uint8) + + # Create an array of z indices + z_indices = np.tile(np.arange(seg_data.shape[2]), (seg_data.shape[0], seg_data.shape[1], 1)) + # Create an array of y indices + y_indices = np.broadcast_to(np.arange(seg_data.shape[1])[..., np.newaxis], seg_data.shape) + + if not from_bottom: + # Cut at the z of the most anteior voxel the lowest vertebrae in the image + last_vert = seg_data[(18 <= seg_data) & (seg_data <= 41)].min() + zmin = -1 + # Get the z - loop to fine the most inferior possible z with spinal canal (label 201). + while zmin == -1 or 201 not in seg_data[..., zmin]: + zmin = z_indices[(seg_data == last_vert) & (y_indices == y_indices[seg_data == last_vert].max())].min() + last_vert += 1 + run_command(f'sct_crop_image -i {temp_path}/img.nii.gz -zmin {zmin} -o {temp_path}/img.nii.gz') + run_command(f'sct_crop_image -i {temp_path}/seg.nii.gz -zmin {zmin} -o {temp_path}/seg.nii.gz') + elif 207 in seg_data: + # Cut at the lowest voxel of T12-L1 IVD + zmax = z_indices[seg_data == 207].min() + run_command(f'sct_crop_image -i {temp_path}/img.nii.gz -zmax {zmax} -o {temp_path}/img.nii.gz') + run_command(f'sct_crop_image -i {temp_path}/seg.nii.gz -zmax {zmax} -o {temp_path}/seg.nii.gz') + + # Copy files from tmp to output destination + shutil.copy(str(temp_path / 'img.nii.gz'), str(output_image_path)) + shutil.copy(str(temp_path / 'seg.nii.gz'), str(output_seg_path)) + finally: + shutil.rmtree(temp_path) + +def run_command(command): + result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + # print(result.stdout) + # print(result.stderr) + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/totalsegmri/utils/generate_labels_sequential.py b/src/totalsegmri/utils/generate_labels_sequential.py new file mode 100644 index 0000000..fa2d727 --- /dev/null +++ b/src/totalsegmri/utils/generate_labels_sequential.py @@ -0,0 +1,290 @@ +import sys, argparse, textwrap +from scipy.ndimage import label, binary_dilation, generate_binary_structure, iterate_structure +from pathlib import Path +import numpy as np +import nibabel as nib +import multiprocessing as mp +from functools import partial +from tqdm.contrib.concurrent import process_map +from totalsegmri.utils.dirpath import DirPath + +def main(): + + # Parse command line arguments + parser = argparse.ArgumentParser( + description=textwrap.dedent(f''' + Label Vertebrae IVD CSF and Spinal Cord from init segmentation. + '''), + epilog=textwrap.dedent(''' + Examples: + generate_labels_sequential -s labels_init -o labels + '''), + formatter_class=argparse.RawTextHelpFormatter + ) + + parser.add_argument( + '--segs-dir', '-s', type=DirPath(), required=True, + help='Folder containing input segmentations.' + ) + parser.add_argument( + '--output-dir', '-o', type=DirPath(True), required=True, + help='Folder to save output segmentations.' + ) + parser.add_argument( + '--subject-dir', '-d', type=str, default=None, nargs='?', const='', + help=textwrap.dedent(''' + Is every subject has its oen direcrory. + If this argument will be provided without value it will look for any directory in the segmentation directory. + If value also provided it will be used as a prefix to subject directory (for example "sub-"), defaults to False (no subjet directory). + '''), + ) + parser.add_argument( + '--subject-subdir', '-u', type=str, default='', + help='Subfolder inside subject folder containing masks (for example "anat"), defaults to no subfolder.' + ) + parser.add_argument( + '--prefix', '-p', type=str, default='', + help='File prefix to work on.' + ) + parser.add_argument( + '--seg-suffix', type=str, default='', + help='Segmentation suffix, defaults to "".' + ) + parser.add_argument( + '--output-seg-suffix', type=str, default='', + help='Suffix for output segmentation, defaults to "".' + ) + parser.add_argument( + '--disc-labels', type=int, nargs='+', default=[], + help='The disc labels.' + ) + parser.add_argument( + '--init-disc', type=lambda x:map(int, x.split(',')), nargs='+', default=[], + help='Init labels list for disc ordered by priority (input_label,output_label !!without space!!). for example 4,224 5,219 6,202' + ) + parser.add_argument( + '--output-disc-step', type=int, default=-1, + help='The step to take between disc labels in the output, defaults to -1.' + ) + parser.add_argument( + '--vertebrea-labels', type=int, nargs='+', default=[], + help='The vertebrae labels.' + ) + parser.add_argument( + '--init-vertebrae', type=lambda x:map(int, x.split(',')), nargs='+', default=[], + help='Init labels list for vertebrae ordered by priority (input_label,output_label !!without space!!). for example 10,41 11,34 12,18' + ) + parser.add_argument( + '--output-vertebrea-step', type=int, default=-1, + help='The step to take between vertebrae labels in the output, defaults to -1.' + ) + parser.add_argument( + '--csf-labels', type=int, nargs='+', default=[], + help='The CSF label.' + ) + parser.add_argument( + '--output-csf-label', type=int, default=201, + help='The CSF label in the output, defaults to 201.' + ) + parser.add_argument( + '--sc-labels', type=int, nargs='+', default=[], + help='The spinal cord label.' + ) + parser.add_argument( + '--output-sc-label', type=int, default=200, + help='The spinal cord label in the output, defaults to 200.' + ) + parser.add_argument( + '--dilation-size', type=int, default=1, + help='Number of voxels to dilate before finding connected voxels to label, defaults to 1 (No dilation).' + ) + parser.add_argument( + '--combine-before-label', action="store_true", default=False, + help='Combine all labels before label continue voxels, defaults to false (label continue voxels separatly for each label).' + ) + parser.add_argument( + '--max-workers', '-w', type=int, default=min(32, mp.cpu_count() + 4), + help='Max worker to run in parallel proccess, defaults to min(32, mp.cpu_count() + 4).' + ) + parser.add_argument( + '--verbose', '-v', type=int, default=1, choices=[0, 1], + help='Verbosity level. 0: Errors/warnings only, 1: Errors/warnings + info (default: 1)' + ) + + try: + args = parser.parse_args() + except BaseException as e: + sys.exit() + + # Get arguments + seg_path = args.seg_dir + output_path = args.output_dir + subject_dir = args.subject_dir + subject_subdir = args.subject_subdir + prefix = args.prefix + seg_suffix = args.seg_suffix + output_seg_suffix = args.output_seg_suffix + disc_labels = args.disc_labels + init_disc = dict(args.init_disc) + output_disc_step = args.output_disc_step + vertebrea_labels = args.vertebrea_labels + init_vertebrae = dict(args.init_vertebrae) + output_vertebrea_step = args.output_vertebrea_step + csf_labels = args.csf_labels + output_csf_label = args.output_csf_label + sc_labels = args.sc_labels + output_sc_label = args.output_sc_label + dilation_size = args.dilation_size + combine_before_label = args.combine_before_label + max_workers = args.max_workers + verbose = args.verbose + + if verbose: + print(textwrap.dedent(f''' + Running {Path(__file__).stem} with the following params: + seg_dir = "{seg_path}" + output_dir = "{output_path}" + subject_dir = "{subject_dir}" + subject_subdir = "{subject_subdir}" + prefix = "{prefix}" + seg_suffix = "{seg_suffix}" + output_seg_suffix = "{output_seg_suffix}" + disc_labels = "{disc_labels}" + init_disc = "{init_disc}" + output_disc_step = "{output_disc_step}" + vertebrea_labels = "{vertebrea_labels}" + init_vertebrae = "{init_vertebrae}" + output_vertebrea_step = "{output_vertebrea_step}" + csf_labels = "{csf_labels}" + output_csf_label = "{output_csf_label}" + sc_labels = "{sc_labels}" + output_sc_label = "{output_sc_label}" + dilation_size = "{dilation_size}" + combine_before_label = "{combine_before_label}" + max_workers = "{max_workers}" + verbose = {verbose} + ''')) + + glob_pattern = "" + if subject_dir is not None: + glob_pattern += f"{subject_dir}*/" + if len(subject_subdir) > 0: + glob_pattern += f"{subject_subdir}/" + glob_pattern += f'{prefix}*{seg_suffix}.nii.gz' + + # Process the NIfTI image and segmentation files + segs_path_list = list(seg_path.glob(glob_pattern)) + + # Create a partially-applied function with the extra arguments + partial_generate_labels_sequential = partial( + generate_labels_sequential, + output_path=output_path, + seg_suffix=seg_suffix, + output_seg_suffix=output_seg_suffix, + disc_labels=disc_labels, + output_disc_step=output_disc_step, + init_disc=init_disc, + vertebrea_labels=vertebrea_labels, + init_vertebrae=init_vertebrae, + output_vertebrea_step=output_vertebrea_step, + csf_labels=csf_labels, + output_csf_label=output_csf_label, + sc_labels=sc_labels, + output_sc_label=output_sc_label, + dilation_size=dilation_size, + combine_before_label=combine_before_label, + ) + + with mp.Pool() as pool: + process_map(partial_generate_labels_sequential, segs_path_list, max_workers=max_workers) + + +def generate_labels_sequential( + seg_path, + output_path, + seg_suffix, + output_seg_suffix, + disc_labels, + init_disc, + output_disc_step, + vertebrea_labels, + init_vertebrae, + output_vertebrea_step, + csf_labels, + output_csf_label, + sc_labels, + output_sc_label, + dilation_size, + combine_before_label, + ): + + output_seg_path = output_path / seg_path.name.replace(f'{seg_suffix}.nii.gz', f'{output_seg_suffix}.nii.gz') + + # Load segmentation + seg = nib.load(seg_path) + seg_data = seg.get_fdata() + + # Convert data to uint8 to avoid issues with segmentation IDs + seg_data_src = seg_data.astype(np.uint8) + + seg_data = np.zeros_like(seg_data_src) + + # Set cord label + seg_data[np.isin(seg_data_src, sc_labels)] = output_sc_label + # Set CSF label + seg_data[np.isin(seg_data_src, csf_labels)] = output_csf_label + + # Create an array of z indices + z_indices = np.broadcast_to(np.arange(seg_data_src.shape[2]), seg_data_src.shape) + + binary_dilation_structure = iterate_structure(generate_binary_structure(3, 1), dilation_size) + + for labels, step, init in ( + (disc_labels, output_disc_step, init_disc), + (vertebrea_labels, output_vertebrea_step, init_vertebrae)): + + if len(labels) == 0: + continue + + # Get labeled + if combine_before_label: + mask_labeled, num_labels = label(binary_dilation(np.isin(seg_data_src, labels), binary_dilation_structure), np.ones((3, 3, 3))) + else: + mask_labeled, num_labels = np.zeros_like(seg_data_src), 0 + for l in labels: + mask = seg_data_src == l + tmp_mask_labeled, tmp_num_labels = label(binary_dilation(mask, binary_dilation_structure), np.ones((3, 3, 3))) + tmp_mask_labeled *= mask + if tmp_num_labels > 0: + mask_labeled[tmp_mask_labeled != 0] = tmp_mask_labeled[tmp_mask_labeled != 0] + num_labels + num_labels += tmp_num_labels + + # Get the z-axis index for each label + mask_labeled_z_indexes = [np.max(z_indices[mask_labeled==i]) for i in range(1, num_labels+1)] + + # Sort the labels by their z-index (reversed) + sorted_labels = [x for _,x in sorted(zip(mask_labeled_z_indexes,range(1,num_labels+1)))][::-1] + + first_label = 0 + for k, v in init.items(): + if k in seg_data_src: + first_label = v - step * sorted_labels.index(np.min(mask_labeled[seg_data_src == k])) + break + + if first_label == 0: + print(f"Error: {seg_path}, some initiation label must be in the segmentation (init: {init.keys()})") + return + + for i in range(num_labels): + seg_data[mask_labeled == sorted_labels[i]] = first_label + step * i + + # Create result segmentation + seg = nib.Nifti1Image(seg_data, seg.affine, seg.header) + seg.set_data_dtype(np.uint8) + # Make sure output directory exists + output_seg_path.parent.mkdir(parents=True, exist_ok=True) + # Save mapped segmentation + nib.save(seg, output_seg_path) + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/totalsegmri/utils/generate_largest_labels.py b/src/totalsegmri/utils/generate_largest_labels.py new file mode 100644 index 0000000..2cc1e54 --- /dev/null +++ b/src/totalsegmri/utils/generate_largest_labels.py @@ -0,0 +1,153 @@ +import sys, argparse, textwrap +from scipy.ndimage import label +from pathlib import Path +import numpy as np +import nibabel as nib +import multiprocessing as mp +from functools import partial +from tqdm.contrib.concurrent import process_map +from totalsegmri.utils.dirpath import DirPath + +def main(): + + # Parse command line arguments + parser = argparse.ArgumentParser( + description=textwrap.dedent(f''' + This script processes NIfTI segmentation files, leaving the largest component for each label. + '''), + epilog=textwrap.dedent(''' + Examples: + generate_largest_labels -s labels_init -o labels + '''), + formatter_class=argparse.RawTextHelpFormatter + ) + + parser.add_argument( + '--seg-dir', '-s', type=DirPath(), required=True, + help='Folder containing input segmentations.' + ) + parser.add_argument( + '--output-dir', '-o', type=DirPath(True), required=True, + help='Folder to save output segmentations.' + ) + parser.add_argument( + '--subject-dir', '-d', type=str, default=None, nargs='?', const='', + help=textwrap.dedent(''' + Is every subject has its oen direcrory. + If this argument will be provided without value it will look for any directory in the segmentation directory. + If value also provided it will be used as a prefix to subject directory (for example "sub-"), defaults to False (no subjet directory). + '''), + ) + parser.add_argument( + '--subject-subdir', '-u', type=str, default='', + help='Subfolder inside subject folder containing masks (for example "anat"), defaults to no subfolder.' + ) + parser.add_argument( + '--prefix', '-p', type=str, default='', + help='File prefix to work on.' + ) + parser.add_argument( + '--seg-suffix', type=str, default='', + help='Segmentation suffix, defaults to "".' + ) + parser.add_argument( + '--output-seg-suffix', type=str, default='', + help='Suffix for output segmentation, defaults to "".' + ) + parser.add_argument( + '--max-workers', '-w', type=int, default=min(32, mp.cpu_count() + 4), + help='Max worker to run in parallel proccess, defaults to min(32, mp.cpu_count() + 4).' + ) + parser.add_argument( + '--verbose', '-v', type=int, default=1, choices=[0, 1], + help='Verbosity level. 0: Errors/warnings only, 1: Errors/warnings + info (default: 1)' + ) + + try: + args = parser.parse_args() + except BaseException as e: + sys.exit() + + # Get arguments + seg_path = args.seg_dir + output_path = args.output_dir + subject_dir = args.subject_dir + subject_subdir = args.subject_subdir + prefix = args.prefix + seg_suffix = args.seg_suffix + output_seg_suffix = args.output_seg_suffix + max_workers = args.max_workers + verbose = args.verbose + + if verbose: + print(textwrap.dedent(f''' + Running {Path(__file__).stem} with the following params: + seg_dir = "{seg_path}" + output_dir = "{output_path}" + subject_dir = "{subject_dir}" + subject_subdir = "{subject_subdir}" + prefix = "{prefix}" + seg_suffix = "{seg_suffix}" + output_seg_suffix = "{output_seg_suffix}" + max_workers = "{max_workers}" + verbose = {verbose} + ''')) + + glob_pattern = "" + if subject_dir is not None: + glob_pattern += f"{subject_dir}*/" + if len(subject_subdir) > 0: + glob_pattern += f"{subject_subdir}/" + glob_pattern += f'{prefix}*{seg_suffix}.nii.gz' + + # Process the NIfTI image and segmentation files + segs_path_list = list(seg_path.glob(glob_pattern)) + + # Create a partially-applied function with the extra arguments + partial_generate_largest_labels = partial( + generate_largest_labels, + output_path=output_path, + seg_suffix=seg_suffix, + output_seg_suffix=output_seg_suffix, + ) + + with mp.Pool() as pool: + process_map(partial_generate_largest_labels, segs_path_list, max_workers=max_workers) + + +def generate_largest_labels( + seg_path, + output_path, + seg_suffix, + output_seg_suffix, + ): + + output_seg_path = output_path / seg_path.name.replace(f'{seg_suffix}.nii.gz', f'{output_seg_suffix}.nii.gz') + + # Load segmentation + seg = nib.load(seg_path) + seg_data = seg.get_fdata() + + # Convert data to uint8 to avoid issues with segmentation IDs + seg_data_src = seg_data.astype(np.uint8) + + seg_data = np.zeros_like(seg_data_src) + + for l in [_ for _ in np.unique(seg_data_src) if _ != 0]: + mask = seg_data_src == l + mask_labeled, num_labels = label(mask, np.ones((3, 3, 3))) + # Find the label of the largest component + label_sizes = np.bincount(mask_labeled.ravel())[1:] # Skip 0 label size + largest_label = label_sizes.argmax() + 1 # +1 because bincount labels start at 0 + seg_data[mask_labeled == largest_label] = l + + # Create result segmentation + seg = nib.Nifti1Image(seg_data, seg.affine, seg.header) + seg.set_data_dtype(np.uint8) + # Make sure output directory exists + output_seg_path.parent.mkdir(parents=True, exist_ok=True) + # Save mapped segmentation + nib.save(seg, output_seg_path) + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/totalsegmri/utils/generate_seg_jpg.py b/src/totalsegmri/utils/generate_seg_jpg.py new file mode 100644 index 0000000..0d2198b --- /dev/null +++ b/src/totalsegmri/utils/generate_seg_jpg.py @@ -0,0 +1,236 @@ +import sys, argparse, textwrap +import numpy as np +import nibabel as nib +from nibabel import freesurfer +from PIL import Image +from nilearn import image as nl_image +import multiprocessing as mp +from functools import partial +from pathlib import Path +from tqdm.contrib.concurrent import process_map +from totalsegmri.utils.dirpath import DirPath +import warnings + +warnings.filterwarnings("ignore") + +def main(): + + # Description and arguments + parser = argparse.ArgumentParser( + description=textwrap.dedent(f''' + This script processes NIfTI (Neuroimaging Informatics Technology Initiative) image and segmentation files. + It combines the specified slice of the image and segmentation files and saves the result as a JPG image + in the specified output folder.''' + ), + formatter_class=argparse.RawTextHelpFormatter + ) + parser.add_argument( + '--images-dir', '-i', type=DirPath(), required=True, + help='The folder where input NIfTI images files are located (required).' + ) + parser.add_argument( + '--segs-dir', '-s', type=DirPath(), required=True, + help='The folder where input NIfTI segmentation files are located (required).' + ) + parser.add_argument( + '--output-dir', '-o', type=DirPath(True), required=True, + help='The folder where output combined JPG images will be saved (required).' + ) + parser.add_argument( + '--subject-dir', '-d', type=str, default=None, nargs='?', const='', + help=textwrap.dedent(''' + Is every subject has its oen direcrory. + If this argument will be provided without value it will look for any directory in the segmentation directory. + If value also provided it will be used as a prefix to subject directory, defaults to False (no subjet directory). + '''), + ) + parser.add_argument( + '--subject-subdir', '-u', type=str, default='', + help='Subfolder inside subject folder containing masks, defaults to no subfolder.' + ) + parser.add_argument( + '--prefix', '-p', type=str, default='', + help='File prefix to work on.' + ) + parser.add_argument( + '--seg-suffix', type=str, default='', + help='Segmentation suffix, defaults to "".' + ) + parser.add_argument( + '--image-suffix', type=str, default='_0000', + help='Image suffix, defaults to "_0000".' + ) + parser.add_argument( + '--orient', '-t', type=str, choices=['sag', 'ax', 'cor'], default='sag', + help='Orientation of the output slice (sagittal, axial, or coronal). Default is "sag".' + ) + parser.add_argument( + '--sliceloc', '-l', type=float, default=0.5, + help='Slice location within the specified orientation (0-1). Default is 0.5 (middle slice).' + ) + parser.add_argument( + '--override', '-r', type=int, default=0, choices=[0, 1], + help='Override existing output files. 0: Do not override, 1: Override. Default is 0.' + ) + parser.add_argument( + '--max-workers', '-w', type=int, default=min(32, mp.cpu_count() + 4), + help='Max worker to run in parallel proccess, defaults to min(32, mp.cpu_count() + 4).' + ) + parser.add_argument( + '--verbose', '-v', type=int, default=1, choices=[0, 1], + help='verbose. 0: Display only errors/warnings, 1: Errors/warnings + info messages. Default is 1.' + ) + + # Parse the command-line arguments + try: + args = parser.parse_args() + except BaseException as e: + sys.exit() + + # Get the command-line argument values + images_path = args.images_dir + segs_path = args.segs_dir + output_path = args.output_dir + subject_dir = args.subject_dir + subject_subdir = args.subject_subdir + prefix = args.prefix + seg_suffix = args.seg_suffix + image_suffix = args.image_suffix + orient = args.orient + sliceloc = args.sliceloc + override = args.override + max_workers = args.max_workers + verbose = args.verbose + + # Print the argument values if verbose is enabled + if verbose: + print(textwrap.dedent(f''' + Running {Path(__file__).stem} with the following params: + images_path = "{images_path}" + segs_path = "{segs_path}" + output_dir = "{output_path}" + subject_dir = "{subject_dir}" + subject_subdir = "{subject_subdir}" + prefix = "{prefix}" + seg_suffix = "{seg_suffix}" + image_suffix = "{image_suffix}" + orient = "{orient}" + sliceloc = "{sliceloc}" + override = "{override}" + max_workers = "{max_workers}" + verbose = "{verbose}" + ''')) + + glob_pattern = "" + if subject_dir is not None: + glob_pattern += f"{subject_dir}*/" + if len(subject_subdir) > 0: + glob_pattern += f"{subject_subdir}/" + glob_pattern += f'{prefix}*{image_suffix}.nii.gz' + + # Process the NIfTI image and segmentation files + images_path_list = list(images_path.glob(glob_pattern)) + + # Create a partially-applied function with the extra arguments + partial_generate_seg_jpg = partial( + generate_seg_jpg, + segs_path=segs_path, + images_path=images_path, + output_path=output_path, + seg_suffix=seg_suffix, + image_suffix=image_suffix, + orient=orient, + sliceloc=sliceloc, + override=override, + ) + + with mp.Pool() as pool: + process_map(partial_generate_seg_jpg, images_path_list, max_workers=max_workers) + + +def generate_seg_jpg(image_path, segs_path, images_path, output_path, seg_suffix, image_suffix, orient, sliceloc, override): + + seg_path = segs_path / image_path.relative_to(images_path).parent / image_path.name.replace(f'{image_suffix}.nii.gz', f'{seg_suffix}.nii.gz') + if not seg_path.is_file(): + seg_path = seg_path.parent / seg_path.name.replace('.nii.gz', '.mgz') + output_jpg_path = output_path / image_path.name.replace(f'{image_suffix}.nii.gz', f'_{orient}_{sliceloc}.jpg') + + if not image_path.is_file() or not seg_path.is_file(): + return + + # Check if the output file exists and if the override flag is set to 0 + if output_jpg_path.exists() and not override: + return + + # Set the seed to ensure consistent colors for each label + np.random.seed(42) + + # Load the NIfTI image file + img = nib.load(image_path) + + # Load the segmentation file (NIfTI or MGZ) + if seg_path.suffix == '.mgz': + seg = freesurfer.load(seg_path) + # Convert the MGHImage to a NIfTI image + seg = nib.Nifti1Image(seg.get_fdata(), seg.affine) + else: + seg = nib.load(seg_path) + + # Reorient the input image to RAS + img = nib.as_closest_canonical(img) + seg = nib.as_closest_canonical(seg) + + # Resample the segmentation to 1X1 mm + seg = nl_image.resample_img(seg, np.eye(3), interpolation='nearest') + + # Resample the image to the same space as the segmentation + img = nl_image.resample_to_img(img, seg) + + # Convert the image and segmentation data into numpy arrays + img_data = img.get_fdata() + seg_data = seg.get_fdata() + + # Find the specified slice + axis={'sag': 0, 'cor': 1, 'ax': 2}[orient] + slice_index = int(sliceloc * img_data.shape[axis]) + slice_img = img_data.take(slice_index, axis=axis) + slice_seg = seg_data.take(slice_index, axis=axis) + + # Generate consistent random colors for each label + unique_labels = np.unique(slice_seg).astype(int) + colors = {} + for label in unique_labels: + np.random.seed(label) + colors[label] = np.random.randint(0, 255, 3) + + # Create a blank color image with the same dimensions as the input image + color_img = np.zeros((*slice_img.shape, 3), dtype=np.uint8) + + # Apply the segmentation mask to the image and assign colors + for label, color in colors.items(): + if label != 0: # Ignore the background (label 0) + mask = slice_seg == label + color_img[mask] = color + + # Normalize the slice to the range 0-255 + normalized_slice = (255 * (slice_img - np.min(slice_img)) / (np.max(slice_img) - np.min(slice_img))).astype(np.uint8) + + # Combine the original grayscale image with the colored segmentation mask + grayscale_img = np.repeat(normalized_slice[:, :, np.newaxis], 3, axis=2).astype(np.uint8) + combined_img = np.where(color_img > 0, color_img, grayscale_img) + + # Rotate the image 90 degrees counter-clockwise and flip it vertically + output_img = np.flipud(combined_img) + output_img = np.rot90(output_img, k=1) + + # Create an Image object from the output image + jpg_image = Image.fromarray(output_img, mode="RGB") + + # Make sure output directory exists + output_jpg_path.parent.mkdir(parents=True, exist_ok=True) + + # Save the Image object as a JPG file + jpg_image.save(output_jpg_path) + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/totalsegmri/utils/map_labels.py b/src/totalsegmri/utils/map_labels.py new file mode 100644 index 0000000..3027aa9 --- /dev/null +++ b/src/totalsegmri/utils/map_labels.py @@ -0,0 +1,176 @@ +import sys, argparse, textwrap, json +import numpy as np +import nibabel as nib +import multiprocessing as mp +from functools import partial +from pathlib import Path +from tqdm.contrib.concurrent import process_map +from totalsegmri.utils.dirpath import DirPath + +def main(): + + # Parse command line arguments + parser = argparse.ArgumentParser( + description=textwrap.dedent(f''' + Map segmentation labels to other labels using json mapping file. + '''), + epilog=textwrap.dedent(''' + Examples: + map_labels -s labels -o labels_maped -m map.json + map_labels -s labels -o labels_maped -m map.json -d sub- -s anat -w 32 + '''), + formatter_class=argparse.RawTextHelpFormatter + ) + + parser.add_argument( + '--segs-dir', '-s', type=DirPath(), required=True, + help='Folder containing input segmentations for each subject.' + ) + parser.add_argument( + '--output-dir', '-o', type=DirPath(True), required=True, + help='Folder to save output segmentations for each subject.' + ) + parser.add_argument( + '--map', '-m', type=argparse.FileType('r', encoding='utf-8'), required=True, + help='JSON file mapping each mask to a unique number, e.g. {"1": 2, "2": 15}' + ) + parser.add_argument( + '--subject-dir', '-d', type=str, default=None, nargs='?', const='', + help=textwrap.dedent(''' + Is every subject has its oen direcrory. + If this argument will be provided without value it will look for any directory in the segmentation directory. + If value also provided it will be used as a prefix to subject directory, defaults to False (no subjet directory). + '''), + ) + parser.add_argument( + '--subject-subdir', '-u', type=str, default='', + help='Subfolder inside subject folder containing masks, defaults to no subfolder.' + ) + parser.add_argument( + '--prefix', '-p', type=str, default='', + help='File prefix to work on.' + ) + parser.add_argument( + '--seg-suffix', type=str, default='', + help='Suffix for output segmentation, defaults to "".' + ) + parser.add_argument( + '--output-seg-suffix', type=str, default='', + help='Suffix for output segmentation, defaults to "".' + ) + parser.add_argument( + '--default-input', action="store_true", default=False, + help='Init output from input, defaults to false (init all to 0).' + ) + parser.add_argument( + '--add-output', action="store_true", default=False, + help='Add new mapped labels to output if the output file exist, defaults to false (Overrite the output file).' + ) + parser.add_argument( + '--max-workers', '-w', type=int, default=min(32, mp.cpu_count() + 4), + help='Max worker to run in parallel proccess, defaults to min(32, mp.cpu_count() + 4).' + ) + parser.add_argument( + '--verbose', '-v', type=int, default=1, choices=[0, 1], + help='Verbosity level. 0: Errors/warnings only, 1: Errors/warnings + info (default: 1)' + ) + + try: + args = parser.parse_args() + except BaseException as e: + sys.exit() + + # Get arguments + seg_path = args.seg_dir + output_path = args.output_dir + map_file = args.map + subject_dir = args.subject_dir + subject_subdir = args.subject_subdir + prefix = args.prefix + seg_suffix = args.seg_suffix + output_seg_suffix = args.output_seg_suffix + default_input = args.default_input + add_output = args.add_output + max_workers = args.max_workers + verbose = args.verbose + + if verbose: + print(textwrap.dedent(f''' + Running {Path(__file__).stem} with the following params: + seg_dir = "{seg_path}" + output_dir = "{output_path}" + map = "{map_file.name}" + subject_dir = "{subject_dir}" + subject_subdir = "{subject_subdir}" + prefix = "{prefix}" + seg_suffix = "{seg_suffix}" + output_seg_suffix = "{output_seg_suffix}" + default_input = "{default_input}" + add_output = "{add_output}" + max_workers = "{max_workers}" + verbose = {verbose} + ''')) + + # Load label mappings from JSON file + map_dict = json.load(map_file) + map_file.close() + + glob_pattern = "" + if subject_dir is not None: + glob_pattern += f"{subject_dir}*/" + if len(subject_subdir) > 0: + glob_pattern += f"{subject_subdir}/" + glob_pattern += f'{prefix}*{seg_suffix}.nii.gz' + + # Process the NIfTI image and segmentation files + segs_path_list = list(seg_path.glob(glob_pattern)) + + # Create a partially-applied function with the extra arguments + partial_map_seg = partial(map_seg, map_dict=map_dict, output_path=output_path, seg_suffix=seg_suffix, output_seg_suffix=output_seg_suffix, add_output=add_output, default_input=default_input) + + with mp.Pool() as pool: + process_map(partial_map_seg, segs_path_list, max_workers=max_workers) + + +def map_seg(seg_path, map_dict, output_path, seg_suffix, output_seg_suffix, add_output, default_input): + + output_seg_path = output_path / seg_path.name.replace(f'{seg_suffix}.nii.gz', f'{output_seg_suffix}.nii.gz') + + # Load segmentation + seg = nib.load(seg_path) + seg_data = seg.get_fdata() + + # Convert data to uint8 to avoid issues with segmentation IDs + seg_data = seg_data.astype(np.uint8) + + # Apply label mapping + mapped_seg_data = np.copy(seg_data) if default_input else np.zeros_like(seg_data) + + # Apply label mapping for all labels that are not mapped to 0 + for orig, new in map_dict.items(): + mapped_seg_data[seg_data==int(orig)] = int(new) + + # Add new mapped labels to output + if add_output and output_seg_path.is_file(): + # Load segmentation + output_seg = nib.load(output_seg_path) + output_seg_data = output_seg.get_fdata() + + # Convert data to uint8 to avoid issues with segmentation IDs + output_seg_data = output_seg_data.astype(np.uint8) + + # Update output from existing output file + mapped_seg_data[output_seg_data != 0] = output_seg_data[output_seg_data != 0] + + # Create result segmentation + mapped_seg = nib.Nifti1Image(mapped_seg_data, seg.affine, seg.header) + mapped_seg.set_data_dtype(np.uint8) + + # Make sure output directory exists + output_seg_path.parent.mkdir(parents=True, exist_ok=True) + + # Save mapped segmentation + nib.save(mapped_seg, output_seg_path) + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/totalsegmri/utils/mha2nii.py b/src/totalsegmri/utils/mha2nii.py new file mode 100644 index 0000000..fbaa58f --- /dev/null +++ b/src/totalsegmri/utils/mha2nii.py @@ -0,0 +1,111 @@ +import sys, argparse, textwrap +import multiprocessing as mp +from pathlib import Path +from tqdm.contrib.concurrent import process_map +from functools import partial +import SimpleITK as sitk +from totalsegmri.utils.dirpath import DirPath + + +def main(): + + # Description and arguments + parser = argparse.ArgumentParser( + description=textwrap.dedent(f''' + This script processes .mha images and convert them into .nii.gz images.''' + ), + formatter_class=argparse.RawTextHelpFormatter + ) + parser.add_argument( + '--images-dir', '-i', type=DirPath(), required=True, + help='The folder where input .mha images files are located (required).' + ) + parser.add_argument( + '--output-dir', '-o', type=DirPath(True), required=True, + help='The folder where output .nii.gz images will be saved (required).' + ) + parser.add_argument( + '--subject-dir', '-d', type=str, default=None, nargs='?', const='', + help=textwrap.dedent(''' + Is every subject has its oen direcrory. + If this argument will be provided without value it will look for any directory in the segmentation directory. + If value also provided it will be used as a prefix to subject directory, defaults to False (no subjet directory). + '''), + ) + parser.add_argument( + '--subject-subdir', '-u', type=str, default='', + help='Subfolder inside subject folder containing masks, defaults to no subfolder.' + ) + parser.add_argument( + '--prefix', '-p', type=str, default='', + help='File prefix to work on.' + ) + parser.add_argument( + '--image-suffix', type=str, default='', + help='Image suffix, defaults to "".' + ) + parser.add_argument( + '--max-workers', '-w', type=int, default=min(32, mp.cpu_count() + 4), + help='Max worker to run in parallel proccess, defaults to min(32, mp.cpu_count() + 4).' + ) + parser.add_argument( + '--verbose', '-v', type=int, default=1, choices=[0, 1], + help='verbose. 0: Display only errors/warnings, 1: Errors/warnings + info messages. Default is 1.' + ) + + # Parse the command-line arguments + try: + args = parser.parse_args() + except BaseException as e: + sys.exit() + + # Get the command-line argument values + images_path = args.images_dir + output_path = args.output_dir + subject_dir = args.subject_dir + subject_subdir = args.subject_subdir + prefix = args.prefix + max_workers = args.max_workers + verbose = args.verbose + + # Print the argument values if verbose is enabled + if verbose: + print(textwrap.dedent(f''' + Running {Path(__file__).stem} with the following params: + images_path = "{images_path}" + output_dir = "{output_path}" + subject_dir = "{subject_dir}" + subject_subdir = "{subject_subdir}" + prefix = "{prefix}" + max_workers = "{max_workers}" + verbose = "{verbose}" + ''')) + + glob_pattern = "" + if subject_dir is not None: + glob_pattern += f"{subject_dir}*/" + if len(subject_subdir) > 0: + glob_pattern += f"{subject_subdir}/" + glob_pattern += f'{prefix}*.mha' + + # Process the NIfTI image and segmentation files + images_path_list = list(images_path.glob(glob_pattern)) + + # Create a partially-applied function with the extra arguments + partial_mha2nii = partial( + mha2nii, + images_path=images_path, + output_path=output_path, + ) + + with mp.Pool() as pool: + process_map(partial_mha2nii, images_path_list, max_workers=max_workers) + + +def mha2nii(image_path, images_path, output_path): + + output_nii_path = output_path / image_path.name.replace(f'.mha', f'.nii.gz') + sitk.WriteImage(sitk.ReadImage(image_path), output_nii_path) + +if __name__ == '__main__': + main() diff --git a/test/script1.sh b/test/script1.sh deleted file mode 100755 index 1504b22..0000000 --- a/test/script1.sh +++ /dev/null @@ -1,2 +0,0 @@ -docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -v "$PWD":/mnt -v /mounts/auto:/mounts/auto:shared -v /data:/data:shared -v /mounts/auto/arman7/totalSegmentatorFolder/Totalsegmentator_dataset/:/Totalsegmentator_dataset/ -v /mounts/auto/ipmsa-results:/data/ipmsa-results:shared -v /mounts/auto/arman7/totalSegmentatorFolder/TotalSegmentatorMRI_SynthSeg:/opt/TotalSegmentatorMRI_SynthSeg -u $(id -u ${USER}):$(id -g ${USER}) --network=host --workdir /mnt --rm -ti armaneshaghi/totalsegmentator-mri:latest python /opt/totalsegmentator-mri/scripts/combine_masks.py -d /opt/TotalSegmentatorMRI_SynthSeg/data/derivatives/manual_masks -o /opt/TotalSegmentatorMRI_SynthSeg/output/MP-RAGE_Masks_Combined -m /opt/totalsegmentator-mri/resources/labels.json -#python totalsegmentator-mri/scripts/combine_masks.py -d TotalSegmentatorMRI_SynthSeg/data/derivatives/manual_masks -o TotalSegmentatorMRI_SynthSeg/output/MP-RAGE_Masks_Combined -m totalsegmentator-mri/resources/labels.json diff --git a/test/script2.sh b/test/script2.sh deleted file mode 100755 index ed78efe..0000000 --- a/test/script2.sh +++ /dev/null @@ -1,6 +0,0 @@ -command="python /opt/totalsegmentator-mri/scripts/build_intensity_stats.py -d /opt/TotalSegmentatorMRI_SynthSeg/data -s /opt/TotalSegmentatorMRI_SynthSeg/output/MP-RAGE_Masks_Combined -o /opt/TotalSegmentatorMRI_SynthSeg/output/MP-RAGE_priors -m /opt/totalsegmentator-mri/resources/labels.json -c /opt/totalsegmentator-mri/resources/classes.json" - -#run on docker - -docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -v "$PWD":/mnt -v /mounts/auto:/mounts/auto:shared -v /data:/data:shared -v /mounts/auto/arman7/totalSegmentatorFolder/Totalsegmentator_dataset/:/Totalsegmentator_dataset/ -v /mounts/auto/ipmsa-results:/data/ipmsa-results:shared -v /mounts/auto/arman7/totalSegmentatorFolder/TotalSegmentatorMRI_SynthSeg:/opt/TotalSegmentatorMRI_SynthSeg -u $(id -u ${USER}):$(id -g ${USER}) --network=host --workdir /mnt --rm -ti armaneshaghi/totalsegmentator-mri:latest ${command} - diff --git a/test/script3.sh b/test/script3.sh deleted file mode 100755 index 4621415..0000000 --- a/test/script3.sh +++ /dev/null @@ -1,6 +0,0 @@ -command="python /opt/totalsegmentator-mri/scripts/combine_masks.py -d /opt/TotalSegmentatorMRI_SynthSeg/Totalsegmentator_dataset -o /opt/TotalSegmentatorMRI_SynthSeg/output/TotalSegmentator_Masks_Combined -m /opt/totalsegmentator-mri/resources/labels.json --subject-prefix s --subject-subdir segmentations --seg-suffix _ct_seg" - -#run on docker - -docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -v "$PWD":/mnt -v /mounts/auto:/mounts/auto:shared -v /data:/data:shared -v /mounts/auto/arman7/totalSegmentatorFolder/Totalsegmentator_dataset/:/Totalsegmentator_dataset/ -v /mounts/auto/ipmsa-results:/data/ipmsa-results:shared -v /mounts/auto/arman7/totalSegmentatorFolder/TotalSegmentatorMRI_SynthSeg:/opt/TotalSegmentatorMRI_SynthSeg -u $(id -u ${USER}):$(id -g ${USER}) --network=host --workdir /mnt --rm -ti armaneshaghi/totalsegmentator-mri:latest ${command} - diff --git a/test/script4.sh b/test/script4.sh deleted file mode 100755 index eb4a3b6..0000000 --- a/test/script4.sh +++ /dev/null @@ -1,4 +0,0 @@ -command="python /opt/totalsegmentator-mri/scripts/generate_image.py -s /opt/TotalSegmentatorMRI_SynthSeg/output/TotalSegmentator_Masks_Combined/sub-0287/anat/sub-0287_ct_seg.nii.gz -p /opt/TotalSegmentatorMRI_SynthSeg/output/MP-RAGE_priors -o /opt/TotalSegmentatorMRI_SynthSeg/output/MP-RAGE_Synthetic/test3" - -#run on docker -docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -v "$PWD":/mnt -v /mounts/auto:/mounts/auto:shared -v /data:/data:shared -v /mounts/auto/arman7/totalSegmentatorFolder/Totalsegmentator_dataset/:/Totalsegmentator_dataset/ -v /mounts/auto/ipmsa-results:/data/ipmsa-results:shared -v /mounts/auto/arman7/totalSegmentatorFolder/TotalSegmentatorMRI_SynthSeg:/opt/TotalSegmentatorMRI_SynthSeg -u $(id -u ${USER}):$(id -g ${USER}) --network=host --workdir /mnt --rm -ti armaneshaghi/totalsegmentator-mri:latest ${command}