├── LICENSE ├── LST_AI ├── __init__.py ├── annotate.py ├── custom_tf.py ├── lst ├── register.py ├── segment.py ├── stats.py ├── strip.py └── utils.py ├── README.md ├── docker ├── Dockerfile └── Readme.md ├── figures └── header.png └── setup.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Julian McGinnis 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /LST_AI/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CompImg/LST-AI/7d5ba1ca1be6f0a1ac7b72a14a247f999c2cc027/LST_AI/__init__.py -------------------------------------------------------------------------------- /LST_AI/annotate.py: -------------------------------------------------------------------------------- 1 | """ 2 | This module contains label information for Output and MSMask. 3 | 4 | Output labels: 5 | 1 == Periventricular 6 | 2 == Juxtacortical 7 | 3 == Subcortical 8 | 4 == Infratentorial 9 | 10 | MSMask labels: 11 | 1 == CSF 12 | 2 == GM 13 | 3 == WM 14 | 4 == Ventricles 15 | 5 == Infratentorial 16 | 17 | """ 18 | import shlex 19 | import subprocess 20 | import os 21 | 22 | import nibabel as nib 23 | import numpy as np 24 | from skimage.measure import label 25 | from skimage.morphology import binary_dilation 26 | 27 | 28 | def annotate_lesions(atlas_t1, atlas_mask, t1w_native, seg_native, out_atlas_warp, 29 | out_atlas_mask_warped, out_annotated_native, n_threads=8): 30 | """ 31 | Annotate lesions in a given image using an atlas. 32 | 33 | Parameters: 34 | ----------- 35 | atlas_t1: str 36 | Path to the atlas T1-weighted image (in MNI space). 37 | atlas_mask: str 38 | Path to the atlas mask image. 39 | t1w_native: str 40 | Path to the T1-weighted image in native space. 41 | seg_native: str 42 | Path to the segmentation image in native space. 43 | out_atlas_warp: str 44 | Path where the warp from the atlas to the patient T1 should be saved. 45 | out_atlas_mask_warped: str 46 | Path where the warped atlas mask should be saved. 47 | out_annotated_native: str 48 | Path where the annotated lesion segmentation in native space should be saved. 49 | 50 | Description: 51 | ------------ 52 | The function performs several tasks: 53 | 1. Registers the atlas to the patient's T1 image using a two-step greedy algorithm. 54 | 2. Warps the atlas mask to the patient's space. 55 | 3. Annotates lesions based on the overlap with the atlas mask. 56 | 4. Saves the annotated segmentation in native space. 57 | 58 | Returns: 59 | -------- 60 | None 61 | 62 | """ 63 | 64 | # Register Atlas -> Patient_T1 using greedy (two-step: rigid first, then deformable) 65 | deformable_call = ( 66 | f"greedy -d 3 -m WNCC 2x2x2 -sv -n 100x50x10" 67 | f" -i {t1w_native} {atlas_t1}" 68 | f" -o {out_atlas_warp}" 69 | f" -threads {n_threads}" 70 | ) 71 | subprocess.run(shlex.split(deformable_call), check=True) 72 | 73 | # Warp MSmask in patient space 74 | warp_call = ( 75 | f"greedy -d 3 -rf {t1w_native} -ri LABEL 0.2vox" 76 | f" -rm {atlas_mask} {out_atlas_mask_warped}" 77 | f" -r {out_atlas_warp}" 78 | f" -threads {n_threads}" 79 | ) 80 | 81 | subprocess.run(shlex.split(warp_call), check=True) 82 | 83 | # Load segmentation and msmask and location-label lesions 84 | seg_nib = nib.load(seg_native) 85 | seg = seg_nib.get_fdata() 86 | seg[seg > 0] = 1 # Make sure seg is binary 87 | msmask = nib.load(out_atlas_mask_warped).get_fdata() 88 | 89 | seg_label = label(seg, connectivity=3) 90 | for lesion_ctr in range(1, seg_label.max() + 1): 91 | # We create a temporary binary mask 92 | # for each lesion & dilate it by 1 93 | # (to "catch" adjacent structures) 94 | temp_mask = np.zeros(seg.shape) 95 | temp_mask[seg_label == lesion_ctr] = 1 96 | temp_mask_dil = binary_dilation( 97 | temp_mask, footprint=np.ones((3, 3, 3))).astype(np.uint8) 98 | 99 | if 4 in msmask[temp_mask_dil == 1]: 100 | seg[temp_mask == 1] = 1 # PV 101 | 102 | elif 2 in msmask[temp_mask_dil == 1]: 103 | seg[temp_mask == 1] = 2 # JC 104 | 105 | elif 5 in msmask[temp_mask_dil == 1]: 106 | seg[temp_mask == 1] = 4 # IT 107 | 108 | else: 109 | seg[temp_mask == 1] = 3 # SC 110 | 111 | # Saving & warping back to T1w space 112 | nib.save(nib.Nifti1Image(seg.astype(np.uint8), 113 | seg_nib.affine, seg_nib.header), 114 | out_annotated_native) 115 | 116 | 117 | if __name__ == "__main__": 118 | 119 | # Only for testing purposes 120 | lst_dir = os.getcwd() 121 | parent_directory = os.path.dirname(lst_dir) 122 | atlas_t1w_path = os.path.join(parent_directory, "atlas", "sub-mni152_space-mni_t1.nii.gz") 123 | atlas_mask_path = os.path.join(parent_directory, "atlas", "sub-mni152_space-mni_msmask.nii.gz") 124 | out_atlas_warp_path = os.path.join(parent_directory, "warp_field.nii.gz") 125 | out_atlas_mask_warped_path = os.path.join(parent_directory, 126 | "tmp_mask.nii.gz") 127 | 128 | # annotate lesion test data 129 | t1w_native_path = os.path.join(parent_directory, "testing", "annotation", 130 | "sub-msseg-test-center01-02_ses-01_space-mni_t1.nii.gz") 131 | seg_native_path = os.path.join(parent_directory, "testing", "annotation", 132 | "sub-msseg-test-center01-02_ses-01_space-mni_seg-manual.nii.gz") 133 | annotated_seg_path = os.path.join(parent_directory, "annotated_segmentation.nii.gz") 134 | 135 | annotate_lesions(atlas_t1=atlas_t1w_path, 136 | atlas_mask=atlas_mask_path, 137 | t1w_native=t1w_native_path, 138 | seg_native=seg_native_path, 139 | out_atlas_warp=out_atlas_warp_path, 140 | out_atlas_mask_warped=out_atlas_mask_warped_path, 141 | out_annotated_native=annotated_seg_path, 142 | n_threads=6) 143 | 144 | # check and remove testing results 145 | gt = os.path.join(parent_directory, "testing", "annotation", 146 | "sub-msseg-test-center01-02_ses-01_space-mni_annotated_seg.nii.gz") 147 | array_gt = nib.load(gt).get_fdata() 148 | array_pred = nib.load(annotated_seg_path).get_fdata() 149 | os.remove(out_atlas_mask_warped_path) 150 | os.remove(out_atlas_warp_path) 151 | os.remove(annotated_seg_path) 152 | np.testing.assert_array_equal(array_gt, array_pred) 153 | -------------------------------------------------------------------------------- /LST_AI/custom_tf.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | def load_custom_model(model_path, compile=False): 5 | """ 6 | Loads a custom TensorFlow Keras model from the specified path. 7 | 8 | This function is specifically designed to handle models that originally used the 9 | `tfa.InstanceNormalization` layer from TensorFlow Addons (tfa). Since tfa is no 10 | longer maintained, this function replaces the `InstanceNormalization` layer with a 11 | custom layer, `CustomGroupNormalization`, to ensure compatibility and avoid the need 12 | for installing tfa. 13 | 14 | Args: 15 | model_path (str): The file path to the saved Keras model. 16 | compile (bool): If True, compiles the model after loading. Defaults to False. 17 | 18 | Returns: 19 | tf.keras.Model: The loaded Keras model with `InstanceNormalization` layers replaced 20 | by `CustomGroupNormalization`. 21 | 22 | Example: 23 | >>> model = load_custom_model('path/to/model.h5', compile=True) 24 | """ 25 | custom_objects = { 26 | 'Addons>InstanceNormalization': CustomGroupNormalization, 27 | } 28 | return tf.keras.models.load_model(model_path, custom_objects=custom_objects, compile=compile) 29 | 30 | 31 | 32 | class CustomGroupNormalization(tf.keras.layers.Layer): 33 | """ 34 | Custom Group Normalization layer for TensorFlow Keras models. 35 | 36 | This class provides an alternative to the `tfa.InstanceNormalization` layer found in 37 | TensorFlow Addons (tfa), which is no longer maintained and not available for MAC ARM platforms. 38 | It facilitates the use of group normalization in models without the dependency on tfa, ensuring 39 | compatibility and broader platform support. 40 | 41 | Args: 42 | groups (int): Number of groups for Group Normalization. Default is -1. 43 | **kwargs: Additional keyword arguments for layer configuration. 44 | """ 45 | def __init__(self, groups=-1, **kwargs): 46 | # Extract necessary arguments from kwargs 47 | self.groups = kwargs.pop('groups', -1) 48 | self.epsilon = kwargs.pop('epsilon', 0.001) 49 | self.center = kwargs.pop('center', True) 50 | self.scale = kwargs.pop('scale', True) 51 | self.beta_initializer = kwargs.pop('beta_initializer', 'zeros') 52 | self.gamma_initializer = kwargs.pop('gamma_initializer', 'ones') 53 | self.beta_regularizer = kwargs.pop('beta_regularizer', None) 54 | self.gamma_regularizer = kwargs.pop('gamma_regularizer', None) 55 | self.beta_constraint = kwargs.pop('beta_constraint', None) 56 | self.gamma_constraint = kwargs.pop('gamma_constraint', None) 57 | 58 | # 'axis' argument is not used in GroupNormalization, so we remove it 59 | kwargs.pop('axis', None) 60 | 61 | super(CustomGroupNormalization, self).__init__(**kwargs) 62 | self.group_norm = tf.keras.layers.GroupNormalization( 63 | groups=self.groups, 64 | epsilon=self.epsilon, 65 | center=self.center, 66 | scale=self.scale, 67 | beta_initializer=self.beta_initializer, 68 | gamma_initializer=self.gamma_initializer, 69 | beta_regularizer=self.beta_regularizer, 70 | gamma_regularizer=self.gamma_regularizer, 71 | beta_constraint=self.beta_constraint, 72 | gamma_constraint=self.gamma_constraint, 73 | **kwargs 74 | ) 75 | 76 | def call(self, inputs, training=None): 77 | return self.group_norm(inputs, training=training) 78 | 79 | def get_config(self): 80 | config = super(CustomGroupNormalization, self).get_config() 81 | config.update({ 82 | 'groups': self.groups, 83 | 'epsilon': self.epsilon, 84 | 'center': self.center, 85 | 'scale': self.scale, 86 | 'beta_initializer': self.beta_initializer, 87 | 'gamma_initializer': self.gamma_initializer, 88 | 'beta_regularizer': self.beta_regularizer, 89 | 'gamma_regularizer': self.gamma_regularizer, 90 | 'beta_constraint': self.beta_constraint, 91 | 'gamma_constraint': self.gamma_constraint 92 | }) 93 | return config -------------------------------------------------------------------------------- /LST_AI/lst: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python -W ignore::DeprecationWarning 2 | # coding: utf-8 3 | 4 | """ 5 | REQUIRES 6 | > greedy 7 | > HD-BET 8 | > MNI atlas files 9 | > Model files 10 | """ 11 | import os 12 | import multiprocessing 13 | import sys 14 | import tempfile 15 | import shutil 16 | import argparse 17 | 18 | from LST_AI.strip import run_hdbet, apply_mask 19 | from LST_AI.register import mni_registration, rigid_reg, apply_warp_label, apply_warp_interp 20 | from LST_AI.segment import unet_segmentation 21 | from LST_AI.annotate import annotate_lesions 22 | from LST_AI.stats import compute_stats 23 | from LST_AI.utils import download_data 24 | 25 | if __name__ == "__main__": 26 | 27 | print("###########################\n") 28 | print("Thank you for using LST-AI. If you publish your results, please cite our paper:") 29 | print("Wiltgen T, McGinnis J, Schlaeger S, Kofler F, Voon C, Berthele A, Bischl D, Grundl L, Will N, Metz M, Schinz D, " 30 | "Sepp D, Prucker P, Schmitz-Koep B, Zimmer C, Menze B, Rueckert D, Hemmer B, Kirschke J, Mühlau M, Wiestler B. " 31 | "LST-AI: A Deep Learning Ensemble for Accurate MS Lesion Segmentation. NeuroImage: Clinical, Volume 42, 2024. " 32 | "https://doi.org/10.1016/j.nicl.2024.103611.") 33 | print("###########################\n") 34 | 35 | parser = argparse.ArgumentParser(description='Segment / Label MS lesions according to McDonald criteria.') 36 | 37 | # Input Images 38 | parser.add_argument('--t1', 39 | dest='t1', 40 | help='Path to T1 image', 41 | type=str, 42 | required=True) 43 | parser.add_argument('--flair', 44 | dest='flair', 45 | help='Path to FLAIR image', 46 | type=str, 47 | required=True) 48 | # Output Images 49 | parser.add_argument('--output', 50 | dest='output', 51 | help='Path to output segmentation nifti path.', 52 | type=str, 53 | required=True) 54 | 55 | parser.add_argument('--existing_seg', 56 | dest='existing_seg', 57 | default='', 58 | help='Path to output segmentation image (will be saved in FLAIR space)', 59 | type=str) 60 | 61 | # Temporary directory 62 | parser.add_argument('--temp', 63 | dest='temp', 64 | default='', 65 | help='Path to temp directory.', 66 | type=str) 67 | 68 | # Mode (segment only, annotate only, segment+annotate (default)) 69 | parser.add_argument('--segment_only', 70 | action='store_true', 71 | dest='segment_only', 72 | help='Only perform the segmentation, and skip lesion annotation.') 73 | 74 | parser.add_argument('--annotate_only', 75 | action='store_true', 76 | dest='annotate_only', 77 | help='Only annotate lesion files without segmentation of lesions.') 78 | 79 | parser.add_argument('--stripped', 80 | action='store_true', 81 | dest='stripped', 82 | help='Images are already skull stripped. Skip skull-stripping.') 83 | 84 | # Model Settings 85 | parser.add_argument('--threshold', 86 | dest='threshold', 87 | help='Threshold for binarizing the joint segmentation (default: 0.5)', 88 | type=float, 89 | default=0.5) 90 | 91 | parser.add_argument('--lesion_threshold', 92 | dest='lesion_threshold', 93 | help='minimum lesion size', 94 | type=int, 95 | default=0) 96 | 97 | parser.add_argument('--clipping', 98 | dest='clipping', 99 | help='Clipping (min & max) for standardization of image intensities (default: 0.5 99.5).', 100 | nargs='+', 101 | type=str, 102 | default=('0.5','99.5')) 103 | 104 | parser.add_argument('--fast-mode', 105 | action='store_true', 106 | dest='fast', 107 | help='Only use one model for hd-bet.') 108 | 109 | parser.add_argument('--probability_map', 110 | action='store_true', 111 | dest='probability_map', 112 | help='Additionally store the probability maps of the three models and of the ensemble network.', 113 | default=False) 114 | 115 | # Computing Resources 116 | parser.add_argument('--device', 117 | dest='device', 118 | help='Either int for GPU ID or "cpu" for CPU (default: 0)', 119 | type=str, 120 | default='0') 121 | 122 | parser.add_argument('--threads', 123 | dest='threads', 124 | help='Number of threads to be used for registration (default: all available)', 125 | type=int, 126 | default=multiprocessing.cpu_count()) 127 | 128 | args = parser.parse_args() 129 | 130 | print(f"Looking for model weights in {os.path.dirname(__file__)}.") 131 | download_data(path=os.path.dirname(__file__)) 132 | 133 | # Sanity Checks 134 | assert os.path.exists(args.t1), 'LST.AI aborted. T1w Image Path does not exist.' 135 | assert os.path.exists(args.flair), 'LST.AI aborted. Flair Image Path does not exist.' 136 | assert str(args.t1).endswith(".nii.gz"), 'Please provide T1w as a zipped nifti.' 137 | assert str(args.flair).endswith(".nii.gz"), 'Please provide FLAIR as a zipped nifti.' 138 | assert not os.path.isfile(args.output), 'Please provide an output path, not a filename.' 139 | assert len(args.clipping)==2, 'Please provide two values for clipping (min & max).' 140 | 141 | # convert input 142 | min_clip, max_clip = float(args.clipping[0]), float(args.clipping[1]) 143 | 144 | if not args.temp: 145 | work_dir = tempfile.mkdtemp(prefix='lst_ai_') 146 | else: 147 | work_dir = os.path.abspath(args.temp) 148 | # make temp directory in case it does not exist 149 | if not os.path.exists(work_dir): 150 | os.makedirs(work_dir) 151 | 152 | # Define Image Paths (original space) 153 | path_org_t1w = os.path.join(work_dir, 'sub-X_ses-Y_space-t1w_T1w.nii.gz') 154 | path_org_flair = os.path.join(work_dir, 'sub-X_ses-Y_space-flair_FLAIR.nii.gz') 155 | path_org_stripped_t1w = os.path.join(work_dir, 'sub-X_ses-Y_space-t1w_desc-stripped_T1w.nii.gz') 156 | path_org_stripped_flair = os.path.join(work_dir, 'sub-X_ses-Y_space-flair_desc-stripped_FLAIR.nii.gz') 157 | 158 | # Define Image Paths (MNI space) 159 | path_mni_t1w = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_T1w.nii.gz') 160 | path_mni_flair = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_FLAIR.nii.gz') 161 | path_mni_stripped_t1w = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_desc-stripped_T1w.nii.gz') 162 | path_mni_stripped_flair = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_desc-stripped_FLAIR.nii.gz') 163 | 164 | # Masks 165 | path_orig_brainmask_t1w = os.path.join(work_dir, 'sub-X_ses-Y_space-t1w_brainmask.nii.gz') 166 | path_orig_brainmask_flair = os.path.join(work_dir, 'sub-X_ses-Y_space-flair_brainmask.nii.gz') 167 | path_mni_brainmask = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_brainmask.nii.gz') 168 | 169 | # Probability Maps MNI & FLAIR 170 | probmap_mni_segmentation = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_seg-lst_prob.nii.gz') 171 | probmap_mni_model1 = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_seg-lst_prob_1.nii.gz') 172 | probmap_mni_model2 = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_seg-lst_prob_2.nii.gz') 173 | probmap_mni_model3 = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_seg-lst_prob_3.nii.gz') 174 | probmap_FLAIR_segmentation = os.path.join(work_dir, 'sub-X_ses-Y_space-FLAIR_seg-lst_prob.nii.gz') 175 | probmap_FLAIR_model1 = os.path.join(work_dir, 'sub-X_ses-Y_space-FLAIR_seg-lst_prob_1.nii.gz') 176 | probmap_FLAIR_model2 = os.path.join(work_dir, 'sub-X_ses-Y_space-FLAIR_seg-lst_prob_2.nii.gz') 177 | probmap_FLAIR_model3 = os.path.join(work_dir, 'sub-X_ses-Y_space-FLAIR_seg-lst_prob_3.nii.gz') 178 | 179 | # Temp Segmentation results 180 | path_orig_segmentation = os.path.join(work_dir, 'sub-X_ses-Y_space-flair_seg-lst.nii.gz') 181 | path_mni_segmentation = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_seg-lst.nii.gz') 182 | path_orig_annotated_segmentation = os.path.join(work_dir, 'sub-X_ses-Y_space-flair_desc-annotated_seg-lst.nii.gz') 183 | path_mni_annotated_segmentation = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_desc-annotated_seg-lst.nii.gz') 184 | 185 | # Output paths (in original space) 186 | filename_output_segmentation = "space-flair_seg-lst.nii.gz" 187 | filename_output_annotated_segmentation = "space-flair_desc-annotated_seg-lst.nii.gz" 188 | 189 | # Stats 190 | filename_output_stats_segmentation = "lesion_stats.csv" 191 | filename_output_stats_annotated_segmentation = "annotated_lesion_stats.csv" 192 | 193 | # affines 194 | path_affine_mni_t1w = os.path.join(work_dir, 'affine_t1w_to_mni.mat') 195 | path_affine_mni_flair = os.path.join(work_dir, 'affine_flair_to_mni.mat') 196 | 197 | # directories 198 | lst_dir = os.path.dirname(__file__) 199 | parent_dir = os.path.dirname(__file__) 200 | 201 | # atlas files 202 | t1w_atlas = os.path.join(parent_dir, "atlas", "sub-mni152_space-mni_t1.nii.gz") 203 | t1w_atlas_stripped = os.path.join(parent_dir, "atlas", "sub-mni152_space-mni_t1bet.nii.gz") 204 | atlas_mask = os.path.join(parent_dir, "atlas", "sub-mni152_space-mni_msmask.nii.gz") 205 | atlas_t1 = t1w_atlas 206 | atlas_warp = os.path.join(work_dir, "atlas_warp_field.nii.gz") 207 | atlas_mask_warped = os.path.join(work_dir,"atlas_mask_warped.nii.gz") 208 | 209 | # model weights 210 | model_directory = os.path.join(parent_dir, 'model') 211 | 212 | # make output path (in case it does not exist) 213 | if not os.path.exists(args.output): 214 | os.makedirs(args.output) 215 | 216 | if os.path.isfile(args.existing_seg) and not args.annotate_only: 217 | print("Existing segmentation may only be used with --annotate_only flag.") 218 | print("Aborted.") 219 | sys.exit(1) 220 | 221 | # Annotation only 222 | if args.annotate_only: 223 | print("LST-AI assumes existing segmentation to be in FLAIR space.") 224 | if os.path.isfile(args.existing_seg): 225 | shutil.copy(args.existing_seg, path_orig_segmentation) 226 | else: 227 | print("Existing segmentation does not exist.") 228 | print("Please provide a valid path via --existing_seg.") 229 | sys.exit(1) 230 | 231 | # Register T1w to FLAIR native 232 | t1w_native= os.path.join(work_dir,"T1w.nii.gz") 233 | native_affine = os.path.join(work_dir,"affine_native_t1_flair.mat") 234 | shutil.copy(args.t1, t1w_native) 235 | 236 | if not args.stripped: 237 | 238 | shutil.copy(args.flair, path_org_flair) 239 | rigid_reg(moving=t1w_native, 240 | fixed=path_org_flair, 241 | affine=native_affine, 242 | destination=path_org_t1w, 243 | n_threads=args.threads) 244 | 245 | # strip only T1w 246 | if args.fast: 247 | run_hdbet(input_image=path_org_t1w, output_image=path_org_stripped_t1w, device=args.device, mode="fast") 248 | else: 249 | run_hdbet(input_image=path_org_t1w, output_image=path_org_stripped_t1w, device=args.device, mode="accurate") 250 | 251 | # move processed mask to correct naming convention 252 | hdbet_mask = path_org_stripped_t1w.replace(".nii.gz", "_mask.nii.gz") 253 | shutil.move(hdbet_mask, path_orig_brainmask_flair) 254 | 255 | # apply brain mask to FLAIR 256 | apply_mask(input_image=path_org_flair, 257 | mask=path_orig_brainmask_flair, 258 | output_image=path_org_stripped_flair) 259 | else: 260 | shutil.copy(args.flair, path_org_stripped_flair) 261 | rigid_reg(moving=t1w_native, 262 | fixed=path_org_stripped_flair, 263 | affine=native_affine, 264 | destination=path_org_stripped_t1w, 265 | n_threads=args.threads) 266 | 267 | 268 | annotate_lesions(atlas_t1=t1w_atlas, 269 | atlas_mask=atlas_mask, 270 | t1w_native=path_org_stripped_t1w, 271 | seg_native=path_orig_segmentation, 272 | out_atlas_warp=atlas_warp, 273 | out_atlas_mask_warped=atlas_mask_warped, 274 | out_annotated_native=path_orig_annotated_segmentation, 275 | n_threads=args.threads) 276 | 277 | shutil.copy(path_orig_annotated_segmentation, 278 | os.path.join(args.output, filename_output_annotated_segmentation)) 279 | 280 | 281 | # Segmentation only + (opt. Annotation) 282 | else: 283 | 284 | if args.stripped: 285 | print("Images have been provided as skull-stripped images. Skipping skull-stripping via HD-BET.") 286 | shutil.copy(args.t1, path_org_stripped_t1w) 287 | shutil.copy(args.flair, path_org_stripped_flair) 288 | 289 | mni_registration(t1w_atlas_stripped, 290 | path_org_stripped_t1w, 291 | path_org_stripped_flair, 292 | path_mni_stripped_t1w, 293 | path_mni_stripped_flair, 294 | path_affine_mni_t1w, 295 | path_affine_mni_flair, 296 | n_threads=args.threads) 297 | 298 | else: 299 | print("Images need to be skull-stripped. Processing with HD-BET.") 300 | shutil.copy(args.t1, path_org_t1w) 301 | shutil.copy(args.flair, path_org_flair) 302 | 303 | # first register to MNI space 304 | mni_registration(t1w_atlas, 305 | path_org_t1w, 306 | path_org_flair, 307 | path_mni_t1w, 308 | path_mni_flair, 309 | path_affine_mni_t1w, 310 | path_affine_mni_flair, 311 | n_threads=args.threads) 312 | 313 | # then skull strip T1w 314 | if args.fast: 315 | run_hdbet(input_image=path_mni_t1w, output_image=path_mni_stripped_t1w, device=args.device, mode="fast") 316 | else: 317 | run_hdbet(input_image=path_mni_t1w, output_image=path_mni_stripped_t1w, device=args.device, mode="accurate") 318 | 319 | # move processed mask to correct naming convention 320 | hdbet_mask = path_mni_stripped_t1w.replace(".nii.gz", "_mask.nii.gz") 321 | shutil.move(hdbet_mask, path_mni_brainmask) 322 | 323 | # then apply brain mask to FLAIR 324 | apply_mask(input_image=path_mni_flair, 325 | mask=path_mni_brainmask, 326 | output_image=path_mni_stripped_flair) 327 | 328 | # create masked images in original space 329 | apply_warp_label(image_org_space=path_org_t1w, 330 | affine=path_affine_mni_t1w, 331 | origin=path_mni_brainmask, 332 | target=path_orig_brainmask_t1w, 333 | reverse=True, 334 | n_threads=args.threads) 335 | 336 | apply_warp_label(image_org_space=path_org_flair, 337 | affine=path_affine_mni_flair, 338 | origin=path_mni_brainmask, 339 | target=path_orig_brainmask_flair, 340 | reverse=True, 341 | n_threads=args.threads) 342 | 343 | apply_mask(input_image=path_org_t1w, 344 | mask=path_orig_brainmask_t1w, 345 | output_image=path_org_stripped_t1w) 346 | 347 | apply_mask(input_image=path_org_flair, 348 | mask=path_orig_brainmask_flair, 349 | output_image=path_org_stripped_flair) 350 | 351 | 352 | # Segmentation 353 | print("Running LST Segmentation.") 354 | unet_segmentation(mni_t1=path_mni_stripped_t1w, 355 | mni_flair=path_mni_stripped_flair, 356 | model_path=model_directory, 357 | output_segmentation_path=path_mni_segmentation, 358 | device=args.device, 359 | probmap=args.probability_map, 360 | output_prob_path=probmap_mni_segmentation, 361 | output_prob1_path=probmap_mni_model1, 362 | output_prob2_path=probmap_mni_model2, 363 | output_prob3_path=probmap_mni_model3, 364 | input_shape=(192,192,192), 365 | threshold=args.threshold, 366 | clipping=(min_clip, max_clip), 367 | lesion_thr=int(args.lesion_threshold)) 368 | 369 | # warp probability maps to FLAIR space if flag is set 370 | if args.probability_map: 371 | apply_warp_interp(image_org_space=path_org_stripped_flair, 372 | affine=path_affine_mni_flair, 373 | origin=probmap_mni_segmentation, 374 | target=probmap_FLAIR_segmentation, 375 | reverse=True, 376 | n_threads=args.threads) 377 | apply_warp_interp(image_org_space=path_org_stripped_flair, 378 | affine=path_affine_mni_flair, 379 | origin=probmap_mni_model1, 380 | target=probmap_FLAIR_model1, 381 | reverse=True, 382 | n_threads=args.threads) 383 | apply_warp_interp(image_org_space=path_org_stripped_flair, 384 | affine=path_affine_mni_flair, 385 | origin=probmap_mni_model2, 386 | target=probmap_FLAIR_model2, 387 | reverse=True, 388 | n_threads=args.threads) 389 | apply_warp_interp(image_org_space=path_org_stripped_flair, 390 | affine=path_affine_mni_flair, 391 | origin=probmap_mni_model3, 392 | target=probmap_FLAIR_model3, 393 | reverse=True, 394 | n_threads=args.threads) 395 | 396 | # warp segmentation back to original FLAIR space 397 | # for now, we do not threshold the lesion sizes of the warped native image 398 | apply_warp_label(image_org_space=path_org_stripped_flair, 399 | affine=path_affine_mni_flair, 400 | target=path_orig_segmentation, 401 | origin=path_mni_segmentation, 402 | reverse=True, 403 | n_threads=args.threads) 404 | 405 | # store the segmentations 406 | shutil.copy(path_orig_segmentation, os.path.join(args.output, filename_output_segmentation)) 407 | 408 | # Annotation 409 | if not args.segment_only: 410 | annotate_lesions(atlas_t1=t1w_atlas, 411 | atlas_mask=atlas_mask, 412 | t1w_native=path_mni_stripped_t1w, 413 | seg_native=path_mni_segmentation, 414 | out_atlas_warp=atlas_warp, 415 | out_atlas_mask_warped=atlas_mask_warped, 416 | out_annotated_native=path_mni_annotated_segmentation, 417 | n_threads=args.threads) 418 | 419 | # warp results back to original FLAIR space 420 | apply_warp_label(image_org_space=path_org_stripped_flair, 421 | affine=path_affine_mni_flair, 422 | origin=path_mni_annotated_segmentation, 423 | target=path_orig_annotated_segmentation, 424 | reverse=True, 425 | n_threads=args.threads) 426 | 427 | # store the segmentations 428 | shutil.copy(path_orig_annotated_segmentation, os.path.join(args.output, filename_output_annotated_segmentation)) 429 | 430 | # Compute Stats of (annotated) segmentation if they exist 431 | if os.path.exists(path_orig_segmentation): 432 | compute_stats(mask_file=path_orig_segmentation, 433 | output_file=os.path.join(args.output, filename_output_stats_segmentation), 434 | multi_class=False) 435 | 436 | if os.path.exists(path_orig_annotated_segmentation): 437 | compute_stats(mask_file=path_orig_annotated_segmentation, 438 | output_file=os.path.join(args.output, filename_output_stats_annotated_segmentation), 439 | multi_class=True) 440 | 441 | print(f"Results in {work_dir}") 442 | if not args.temp: 443 | print("Delete temporary directory: {work_dir}") 444 | shutil.rmtree(work_dir) 445 | print("Done.") 446 | -------------------------------------------------------------------------------- /LST_AI/register.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | import shlex 4 | 5 | def mni_registration(atlas_t1, path_org_t1, path_org_flair, 6 | path_mni_t1, path_mni_flair, 7 | path_t1_affine, path_flair_affine, 8 | n_threads=1): 9 | """ 10 | Perform registration of T1 and FLAIR images into MNI space. 11 | 12 | Parameters: 13 | ---------- 14 | atlas_t1 : str 15 | Path to T1 atlas image. 16 | path_org_t1 : str 17 | Path to original T1 image. 18 | path_org_flair : str 19 | Path to original FLAIR image. 20 | path_mni_t1 : str 21 | Path to MNI T1 image. 22 | path_mni_flair : str 23 | Path to MNI FLAIR image. 24 | path_t1_affine : str 25 | Path to save affine for T1 registration. 26 | path_flair_affine : str 27 | Path to save affine for FLAIR registration. 28 | threads : int, optional 29 | Number of threads to use for registration. Default is 1. 30 | 31 | Returns: 32 | -------- 33 | None 34 | 35 | """ 36 | 37 | # Register T1 -> Atlas_T1 (6 DOF) 38 | rigid_call = (f"greedy -d 3 -a -dof 6 -m NMI -ia-image-centers " 39 | f"-n 100x50x10 -i {atlas_t1} {path_org_t1} -o {path_t1_affine} -threads {n_threads}") 40 | subprocess.run(shlex.split(rigid_call), check=True) 41 | 42 | warp_call = (f"greedy -d 3 -rf {atlas_t1} -rm {path_org_t1} " 43 | f"{path_mni_t1} -r {path_t1_affine} -threads {n_threads}") 44 | subprocess.run(shlex.split(warp_call), check=True) 45 | 46 | # Register FLAIR -> T1 in atlas space and save the affine for later 47 | rigid_call = (f"greedy -d 3 -a -dof 6 -m NMI -ia-image-centers " 48 | f"-n 100x50x10 -i {path_mni_t1} {path_org_flair} -o {path_flair_affine} -threads {n_threads}") 49 | subprocess.run(shlex.split(rigid_call), check=True) 50 | 51 | warp_call = (f"greedy -d 3 -rf {atlas_t1} -rm {path_org_flair} " 52 | f"{path_mni_flair} -r {path_flair_affine} -threads {n_threads}") 53 | subprocess.run(shlex.split(warp_call), check=True) 54 | 55 | def rigid_reg(moving, fixed, affine, destination, n_threads): 56 | """ 57 | Perform rigid registration. 58 | 59 | Parameters: 60 | ---------- 61 | moving : str 62 | Path to moving image. 63 | fixed : str 64 | Path to fixed image. 65 | affine : str 66 | Path to affine. 67 | destination : str 68 | Path to registered image. 69 | n_threads : int 70 | Number of threads used for registration. 71 | """ 72 | # Register moving -> fixed (6 DOF) 73 | rigid_call = (f"greedy -d 3 -a -dof 6 -m NMI -ia-image-centers " 74 | f"-n 100x50x10 -i {fixed} {moving} -o {affine} -threads {n_threads}") 75 | subprocess.run(shlex.split(rigid_call), check=True) 76 | 77 | warp_call = (f"greedy -d 3 -rf {fixed} -rm {moving} " 78 | f"{destination} -r {affine} -threads {n_threads}") 79 | subprocess.run(shlex.split(warp_call), check=True) 80 | 81 | 82 | def apply_warp_label(image_org_space, affine, origin, target, reverse=False, n_threads=1): 83 | """ 84 | Warps an image between its original space and target space. 85 | 86 | Parameters: 87 | image_org_space: str - The image in its original space. 88 | affine: str - The affine transformation file. 89 | origin: str - The origin image file. 90 | target: str - The target image file. 91 | reverse: bool - If True, warps from target space to original space; otherwise, 92 | from original space to target space. 93 | threads : int, optional 94 | Number of threads to use for registration. Default is 1. 95 | """ 96 | if reverse: 97 | warp_call = ( 98 | f"greedy -threads {n_threads} -d 3 -rf {image_org_space} -ri LABEL 0.2vox -rm {origin} " 99 | f"{target} -r {affine},-1" 100 | ) 101 | else: 102 | warp_call = ( 103 | f"greedy -threads {n_threads} -d 3 -rf {image_org_space} -ri LABEL 0.2vox -rm {origin} " 104 | f"{target} -r {affine}" 105 | ) 106 | 107 | subprocess.run(shlex.split(warp_call), check=True) 108 | 109 | def apply_warp_interp(image_org_space, affine, origin, target, reverse=False, n_threads=1): 110 | """ 111 | Warps an image between its original space and target space adn applies linear interpolation. 112 | 113 | Parameters: 114 | image_org_space: str - The image in its original space. 115 | affine: str - The affine transformation file. 116 | origin: str - The origin image file. 117 | target: str - The target image file. 118 | reverse: bool - If True, warps from target space to original space; otherwise, 119 | from original space to target space. 120 | threads : int, optional 121 | Number of threads to use for registration. Default is 1. 122 | """ 123 | if reverse: 124 | warp_call = ( 125 | f"greedy -threads {n_threads} -d 3 -rf {image_org_space} -ri LINEAR -rm {origin} " 126 | f"{target} -r {affine},-1" 127 | ) 128 | else: 129 | warp_call = ( 130 | f"greedy -threads {n_threads} -d 3 -rf {image_org_space} -ri LINEAR -rm {origin} " 131 | f"{target} -r {affine}" 132 | ) 133 | 134 | subprocess.run(shlex.split(warp_call), check=True) 135 | 136 | if __name__ == "__main__": 137 | 138 | # Testing only 139 | 140 | # Working directory 141 | script_dir = os.getcwd() 142 | parent_directory = os.path.dirname(script_dir) 143 | 144 | atlas_t1_path = os.path.join(parent_directory, "atlas", "sub-mni152_space-mni_t1.nii.gz") 145 | atlas_mask_path = os.path.join(parent_directory, "atlas", "sub-mni152_space-mni_msmask.nii.gz") 146 | 147 | # Example data 148 | path_org_t1_path = os.path.join(parent_directory, "testing", "segmentation", "T1W.nii.gz") 149 | path_org_flair_path = os.path.join(parent_directory, "testing", "segmentation", "FLAIR.nii.gz") 150 | 151 | # Temp directory 152 | temp_path = os.path.join(parent_directory, "temp") 153 | # If temp directory does not exist, make it 154 | if not os.path.exists(temp_path): 155 | os.makedirs(temp_path) 156 | 157 | # Where to store the results 158 | path_mni_t1_path = os.path.join(temp_path, "mni_t1.nii.gz") 159 | path_mni_flair_path = os.path.join(temp_path, "mni_flair.nii.gz") 160 | path_t1_affine_path = os.path.join(temp_path, "t1_mni_affine.mat") 161 | path_flair_affine_path = os.path.join(temp_path, "flair_mni_affine.mat") 162 | 163 | mni_registration( 164 | atlas_t1=atlas_t1_path, 165 | path_org_t1=path_org_t1_path, 166 | path_org_flair=path_org_flair_path, 167 | path_mni_t1=path_mni_t1_path, 168 | path_mni_flair=path_mni_flair_path, 169 | path_t1_affine=path_t1_affine_path, 170 | path_flair_affine=path_flair_affine_path 171 | ) 172 | print("Finished!") 173 | -------------------------------------------------------------------------------- /LST_AI/segment.py: -------------------------------------------------------------------------------- 1 | import os 2 | import logging 3 | logging.getLogger('tensorflow').disabled = True 4 | import numpy as np 5 | import nibabel as nib 6 | from scipy.ndimage import label, generate_binary_structure 7 | 8 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 9 | import tensorflow as tf 10 | 11 | from LST_AI.custom_tf import load_custom_model 12 | 13 | 14 | # stolen and adapted from SCT 6.5 math.py line 421 15 | def remove_small_objects(data, dim_lst, unit='mm3', thr=0): 16 | """Removes all unconnected objects smaller than the minimum specified size. 17 | 18 | Adapted from: 19 | https://github.com/ivadomed/ivadomed/blob/master/ivadomed/postprocessing.py#L327 20 | and 21 | https://github.com/ivadomed/ivadomed/blob/master/ivadomed/postprocessing.py#L224 22 | 23 | Args: 24 | data (ndarray): Input data. 25 | dim_lst (list): Dimensions of a voxel in mm. 26 | unit (str): Indicates the units of the objects: "mm3" or "vox" 27 | thr (float): Minimal object size to keep in input data. 28 | 29 | Returns: 30 | ndarray: Array with small objects. 31 | """ 32 | print(f"Thresholding lesions at [mm3]:{np.prod(dim_lst)}") 33 | 34 | px, py, pz = dim_lst 35 | 36 | # structuring element that defines feature connections 37 | bin_structure = generate_binary_structure(3, 2) 38 | 39 | data_label, n = label(data, structure=bin_structure) 40 | 41 | if unit == 'mm3': 42 | size_min = np.round(thr / (px * py * pz)) 43 | else: 44 | print('Please choose a different unit for removeSmall. Choices: vox or mm3') 45 | exit() 46 | 47 | for idx in range(1, n + 1): 48 | data_idx = (data_label == idx).astype(int) 49 | n_nonzero = np.count_nonzero(data_idx) 50 | 51 | if n_nonzero < size_min: 52 | data[data_label == idx] = 0 53 | 54 | return data 55 | 56 | 57 | 58 | def unet_segmentation(model_path, mni_t1, mni_flair, output_segmentation_path, 59 | output_prob_path, output_prob1_path, output_prob2_path, output_prob3_path, 60 | device='cpu', probmap=False, input_shape=(192,192,192), threshold=0.5, clipping=(0.5,99.5), lesion_thr=0): 61 | """ 62 | Segment medical images using ensemble of U-Net models. 63 | 64 | This function utilizes pre-trained U-Net models to perform segmentation 65 | on T1 and FLAIR images in the MNI space. The segmentation output is a 66 | binary lesion mask, which is then saved to a specified path. 67 | 68 | Parameters: 69 | ----------- 70 | model_path : str 71 | Directory path containing the U-Net models. 72 | mni_t1 : str 73 | Path to the T1-weighted image in MNI space. 74 | mni_flair : str 75 | Path to the FLAIR image in MNI space. 76 | output_segmentation_path : str 77 | Path to save the resulting binary segmentation mask. 78 | input_shape : tuple of int, optional 79 | Expected shape of the input images for the U-Net model. 80 | Default is (192, 192, 192). 81 | threshold : float, optional 82 | Segmentation threshold to determine the binary mask from the U-Net's 83 | output. Pixels with values above this threshold in the U-Net output 84 | will be set to 1 in the binary mask, and others to 0. Default is 0.5. 85 | clipping : list of floats, optional 86 | Min and Max values for the np.clip option which applies clipping for 87 | the standardization of image intensities. Default is min=0.5 and max=99.5. 88 | 89 | Returns: 90 | -------- 91 | None 92 | The function saves the binary segmentation mask to the path specified 93 | in `output_segmentation_path`. 94 | 95 | """ 96 | 97 | tf_device = '/CPU:0' if device == 'cpu' else f'/GPU:{device}' 98 | 99 | def adapt_shape(img_arr): 100 | #Crops input image array to target shape; also returns information how to re-zero-pad 101 | difference_0 = img_arr.shape[0] - input_shape[0] 102 | difference_0_l = (difference_0 // 2)+(difference_0 % 2) 103 | difference_0_r = (difference_0 // 2) 104 | 105 | difference_1 = img_arr.shape[1] - input_shape[1] 106 | difference_1_l = (difference_1 // 2)+(difference_1 % 2) 107 | difference_1_r = (difference_1 // 2) 108 | 109 | difference_2 = img_arr.shape[2] - input_shape[2] 110 | difference_2_l = (difference_2 // 2)+(difference_2 % 2) 111 | difference_2_r = (difference_2 // 2) 112 | 113 | img_arr_cropped = img_arr[difference_0_l : img_arr.shape[0] - difference_0_r, difference_1_l : img_arr.shape[1] - difference_1_r,difference_2_l : img_arr.shape[2] - difference_2_r] 114 | 115 | return img_arr_cropped.astype(np.float32), [difference_0_l,difference_0_r,difference_1_l,difference_1_r,difference_2_l,difference_2_r] 116 | 117 | def preprocess_intensities(img_arr, clipping): 118 | #Standardize image intensities to [0;1] 119 | temp_bm = np.zeros(img_arr.shape) 120 | temp_bm[img_arr != 0] = 1 121 | img_arr = np.clip(img_arr, 122 | a_min=np.percentile(img_arr[temp_bm != 0],clipping[0]), 123 | a_max=np.percentile(img_arr[temp_bm != 0],clipping[1]) ) 124 | img_arr -= img_arr[temp_bm == 1].min() 125 | img_arr = img_arr / img_arr[temp_bm == 1].max() 126 | img_arr *= temp_bm 127 | 128 | return img_arr.astype(np.float32) 129 | 130 | # weight files 131 | unet_mdls = [ 132 | "UNet3D_MS_final_mdlA.h5", 133 | "UNet3D_MS_final_mdlB.h5", 134 | "UNet3D_MS_final_mdlC.h5" 135 | ] 136 | unet_mdls = [os.path.join(model_path, x) for x in unet_mdls] 137 | 138 | # Load and preprocess images 139 | t1_nib = nib.load(mni_t1) 140 | t1 = t1_nib.get_fdata() 141 | flair_nib = nib.load(mni_flair) 142 | flair = flair_nib.get_fdata() 143 | 144 | t1, shape_lst = adapt_shape(t1) 145 | flair, _ = adapt_shape(flair) 146 | 147 | t1 = preprocess_intensities(t1, clipping) 148 | flair = preprocess_intensities(flair, clipping) 149 | 150 | joint_seg = np.zeros(t1.shape) 151 | print(f"Running segmentation on {tf_device}.") 152 | 153 | # define list of model probability maps 154 | output_prob_list = [output_prob1_path, output_prob2_path, output_prob3_path] 155 | 156 | for i, model in enumerate(unet_mdls): 157 | with tf.device(tf_device): 158 | print(f"Running model {i}. ") 159 | mdl = load_custom_model(model, compile=False) 160 | 161 | img_image = np.stack([flair, t1], axis=-1) 162 | img_image = np.expand_dims(img_image, axis=0) 163 | with tf.device(tf_device): 164 | out_seg = mdl(img_image)[0] # Will return a len(2) list of [out_seg, out_ds] 165 | out_seg = np.squeeze(out_seg) 166 | 167 | if probmap: 168 | # add zero values to the probability map 169 | out_seg_pad = np.pad(out_seg, 170 | ((shape_lst[0], shape_lst[1]), (shape_lst[2], shape_lst[3]), (shape_lst[4], shape_lst[5])), 171 | 'constant', constant_values=0.) 172 | # save the probability map 173 | nib.save(nib.Nifti1Image(out_seg_pad.astype(np.float32), 174 | flair_nib.affine, 175 | flair_nib.header), 176 | output_prob_list[i]) 177 | 178 | joint_seg += out_seg 179 | 180 | joint_seg /= len(unet_mdls) 181 | 182 | if probmap: 183 | # add zero values to the probability map 184 | joint_seg_pad = np.pad(joint_seg, 185 | ((shape_lst[0], shape_lst[1]), (shape_lst[2], shape_lst[3]), (shape_lst[4], shape_lst[5])), 186 | 'constant', constant_values=0.0) 187 | # save the probability map 188 | nib.save(nib.Nifti1Image(joint_seg_pad.astype(np.float32), 189 | flair_nib.affine, 190 | flair_nib.header), 191 | output_prob_path) 192 | 193 | out_binary = np.zeros(t1.shape) 194 | out_binary[joint_seg > threshold] = 1 195 | 196 | out_binary = np.pad( 197 | out_binary, 198 | ((shape_lst[0], shape_lst[1]), (shape_lst[2], shape_lst[3]), (shape_lst[4], shape_lst[5])), 199 | 'constant', constant_values=0. 200 | ) 201 | 202 | out_binary = remove_small_objects(out_binary, flair_nib.header.get_zooms(), unit="mm3", thr=int(lesion_thr)) 203 | 204 | nib.save(nib.Nifti1Image(out_binary.astype(np.uint8), 205 | flair_nib.affine, 206 | flair_nib.header), 207 | output_segmentation_path) 208 | 209 | 210 | if __name__ == "__main__": 211 | # Testing only 212 | # Working directory 213 | script_dir = os.getcwd() 214 | parent_dir = os.path.dirname(script_dir) 215 | 216 | model_dir = os.path.join(parent_dir, "model") 217 | test_dir = os.path.join(parent_dir, "testing", "annotation") 218 | t1_path = os.path.join(test_dir, "sub-msseg-test-center01-02_ses-01_space-mni_t1.nii.gz") 219 | flair_path = t1_path 220 | output_path = os.path.join(parent_dir, "testing", "seg_mni.nii.gz") 221 | gt_path = os.path.join(test_dir, "sub-msseg-test-center01-02_ses-01_space-mni_seg-unet.nii.gz") 222 | 223 | unet_segmentation(mni_t1=t1_path, 224 | mni_flair=flair_path, 225 | model_path=model_dir, 226 | output_segmentation_path=output_path) 227 | 228 | # Check and remove testing results 229 | gt_data = nib.load(gt_path).get_fdata() 230 | pred_data = nib.load(output_path).get_fdata() 231 | os.remove(output_path) 232 | np.testing.assert_array_equal(gt_data, pred_data) 233 | -------------------------------------------------------------------------------- /LST_AI/stats.py: -------------------------------------------------------------------------------- 1 | import nibabel as nib 2 | import numpy as np 3 | import csv 4 | import argparse 5 | from scipy.ndimage import label 6 | 7 | def compute_stats(mask_file, output_file, multi_class): 8 | """ 9 | Compute statistics from a lesion mask and save the results to a CSV file. 10 | 11 | Parameters: 12 | mask_file (str): Path to the input mask file in NIfTI format. 13 | output_file (str): Path to the output CSV file where results will be saved. 14 | multi_class (bool): Flag indicating whether the mask contains multiple classes (True) or is binary (False). 15 | 16 | This function calculates the number of lesions, the number of voxels in lesions, and the total lesion volume. 17 | If `multi_class` is True, these statistics are calculated for each lesion class separately. 18 | """ 19 | # Load the mask file 20 | mask = nib.load(mask_file) 21 | mask_data = mask.get_fdata() 22 | 23 | # Voxel dimensions to calculate volume 24 | voxel_dims = mask.header.get_zooms() 25 | 26 | results = [] 27 | 28 | if multi_class: 29 | # Multi-class processing 30 | lesion_labels = [1, 2, 3, 4] 31 | label_names = { 32 | 1: 'Periventricular', 33 | 2: 'Juxtacortical', 34 | 3: 'Subcortical', 35 | 4: 'Infratentorial' 36 | } 37 | 38 | for lesion_label in lesion_labels: 39 | class_mask = mask_data == lesion_label 40 | 41 | # Count lesions (connected components) for each class 42 | _ , num_lesions = label(class_mask) 43 | 44 | voxel_count = np.count_nonzero(class_mask) 45 | volume = voxel_count * np.prod(voxel_dims) 46 | 47 | results.append({ 48 | 'Region': label_names[lesion_label], 49 | 'Num_Lesions': num_lesions, 50 | 'Num_Vox': voxel_count, 51 | 'Lesion_Volume': volume 52 | }) 53 | 54 | else: 55 | # Binary mask processing 56 | # Assert that only two unique values are present (0 and 1) 57 | unique_values = np.unique(mask_data) 58 | assert len(unique_values) <= 2, "Binary mask must contain no more than two unique values." 59 | 60 | # Count lesions (connected components) in binary mask 61 | _, num_lesions = label(mask_data > 0) 62 | 63 | voxel_count = np.count_nonzero(mask_data) 64 | volume = voxel_count * np.prod(voxel_dims) 65 | 66 | results.append({ 67 | 'Num_Lesions': num_lesions, 68 | 'Num_Vox': voxel_count, 69 | 'Lesion_Volume': volume 70 | }) 71 | 72 | # Save results to CSV 73 | with open(output_file, 'w', newline='') as file: 74 | writer = csv.writer(file) 75 | if multi_class: 76 | writer.writerow(['Region', 'Num_Lesions', 'Num_Vox', 'Lesion_Volume']) 77 | for result in results: 78 | writer.writerow([result['Region'], result['Num_Lesions'], result['Num_Vox'], result['Lesion_Volume']]) 79 | else: 80 | writer.writerow(['Num_Lesions', 'Num_Vox', 'Lesion_Volume']) 81 | for result in results: 82 | writer.writerow([result['Num_Lesions'], result['Num_Vox'], result['Lesion_Volume']]) 83 | 84 | if __name__ == "__main__": 85 | """ 86 | Main entry point of the script. Parses command-line arguments and calls the compute_stats function. 87 | """ 88 | parser = argparse.ArgumentParser(description='Process a lesion mask file.') 89 | parser.add_argument('--in', dest='input_file', required=True, help='Input mask file path') 90 | parser.add_argument('--out', dest='output_file', required=True, help='Output CSV file path') 91 | parser.add_argument('--multi-class', dest='multi_class', action='store_true', help='Flag for multi-class processing') 92 | 93 | args = parser.parse_args() 94 | 95 | compute_stats(args.input_file, args.output_file, args.multi_class) 96 | -------------------------------------------------------------------------------- /LST_AI/strip.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | import shlex 3 | import nibabel as nib 4 | import numpy as np 5 | 6 | def run_hdbet(input_image, output_image, device, mode="accurate"): 7 | """ 8 | Runs the HD-BET tool to perform brain extraction on an input image. 9 | 10 | Parameters: 11 | input_image (str): Path to the input image file. 12 | output_image (str): Path for the output image file. 13 | device (str): The device to use for computation, either a GPU device number or 'cpu'. 14 | mode (str, optional): Operation mode of HD-BET. Can be 'accurate' or 'fast'. Default is 'accurate'. 15 | 16 | Raises: 17 | AssertionError: If an unknown mode is provided. 18 | 19 | This function utilizes HD-BET, a tool for brain extraction from MRI images. Depending on the chosen mode 20 | and device, it executes the appropriate command. 21 | """ 22 | assert mode in ["accurate","fast"], 'Unknown HD-BET mode. Please choose either "accurate" or "fast"' 23 | 24 | if "cpu" in str(device).lower(): 25 | bet_call = f"hd-bet -i {input_image} -device cpu -mode {mode} -tta 0 -o {output_image}" 26 | else: 27 | bet_call = f"hd-bet -i {input_image} -device {device} -mode accurate -tta 1 -o {output_image}" 28 | 29 | subprocess.run(shlex.split(bet_call), check=True) 30 | 31 | def apply_mask(input_image, mask, output_image): 32 | """ 33 | Applies a mask to an input image and saves the result. 34 | 35 | Parameters: 36 | input_image (str): Path to the input image file. 37 | mask (str): Path to the mask image file. 38 | output_image (str): Path for the output image file where the masked image will be saved. 39 | 40 | This function loads a brain mask and an input image, applies the mask to the input image, 41 | and then saves the result. The mask and the input image are expected to be in a compatible format 42 | and spatial alignment. 43 | """ 44 | brain_mask_arr = nib.load(mask).get_fdata() 45 | image_nib = nib.load(input_image) 46 | image_arr = np.multiply(np.squeeze(image_nib.get_fdata()), np.squeeze(brain_mask_arr)) 47 | nib.save(nib.Nifti1Image(image_arr.astype(np.float32), image_nib.affine, image_nib.header), output_image) 48 | -------------------------------------------------------------------------------- /LST_AI/utils.py: -------------------------------------------------------------------------------- 1 | from io import open 2 | import os 3 | import zipfile 4 | from urllib import request 5 | 6 | def download_data(path): 7 | """ 8 | Downloads required model weights, binaries and atlas files for usage. 9 | """ 10 | url = "https://github.com/CompImg/LST-AI/releases/download/v1.0.0/lst_data.zip" 11 | 12 | target_path = "lst_data.zip" 13 | extract_path = path # This is the base directory. 14 | 15 | atlas_path = os.path.join(extract_path, 'atlas') 16 | binary_path = os.path.join(extract_path, 'binaries') 17 | model_path = os.path.join(extract_path, 'model') 18 | 19 | paths_to_check = [atlas_path, binary_path, model_path] 20 | 21 | # Check if all paths exist. 22 | if not all(os.path.exists(path) for path in paths_to_check): 23 | print("Downloading data...") 24 | # Download the zip file if it doesn't exist. 25 | if not os.path.exists(target_path): 26 | with request.urlopen(url) as response, open(target_path, 'wb') as out_file: 27 | data = response.read() 28 | out_file.write(data) 29 | 30 | # Unzip the file to the base directory. 31 | with zipfile.ZipFile(target_path, 'r') as zip_ref: 32 | zip_ref.extractall(extract_path) 33 | 34 | # Remove the ZIP file after extracting its contents. 35 | os.remove(target_path) 36 | print("Completed.") 37 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # LST-AI - Deep Learning Ensemble for Accurate MS Lesion Segmentation 2 | 3 | Welcome to our codebase for LST-AI, the deep learning-based successor of the original [Lesion Segmentation Toolbox (LST)](https://www.applied-statistics.de/lst.html) by [Schmidt et al.](https://www.sciencedirect.com/science/article/abs/pii/S1053811911013139) 4 | LST-AI was collaboratively developed by the Department of Neurology and Department of Neuroradiology, Klinikum rechts der Isar at the Technical University of Munich, and the Department of Computer Science at the Technical University of Munich. 5 | 6 | Overview 7 | 8 | 9 | **Disclaimer:** LST-AI is a research-only tool for MS Lesion Segmentation and has not been validated, licensed or approved for any clinical usage. 10 | 11 | ## What is different, why or when should I switch?! 12 | * LST-AI is an advanced deep learning-based extension of the original LST with improved performance and additional features. 13 | * LST-AI constitutes a completely new framework and has been developed from scratch. 14 | * While LST depends on MATLAB, we offer LST-AI as a python-based tool which makes it available to the whole community. 15 | * We suggest using LST or LST-AI according to your type of data: 16 | * A 3D T1-weighted and a 3D FLAIR sequence are available: (new) LST-AI 17 | * Only a 3D FLAIR sequence is available: (old) [LST](https://www.applied-statistics.de/lst.html) (LPA) 18 | * Only a 3D T1-weighted sequence is available: not covered by any LST(-AI) version 19 | * If a 3D T1-weighted and a non-3D FLAIR sequence are available, please try both (new) LST-AI or (old) [LST](https://www.applied-statistics.de/lst.html) (LGA or LPA) 20 | 21 | 22 | ## Usage 23 | To allow the usage of LST-AI on different platforms and online/offline usage, we provide LST-AI as a Python package and Docker (CPU and GPU-Docker versions available). 24 | 25 | ### Installing the Python package 26 | 27 | LST-AI is a Python-based package. For Debian-based systems, you can install all required packages via `apt`: 28 | 29 | ``` 30 | apt-get update && apt-get install -y \ 31 | git \ 32 | wget \ 33 | unzip \ 34 | python3 \ 35 | python3-pip 36 | ``` 37 | 38 | Under the hood, LST also wraps [HD-BET](https://github.com/MIC-DKFZ/HD-BET) and [greedy](https://github.com/pyushkevich/greedy). 39 | We guide you through the download/compilation for greedy and installation for HD-BET in the following process. If you encounter specific issues with these packages, let us know in an issue and/or consult the GitHub repositories. 40 | 41 | 1. Make a new directory for LST-AI 42 | ```bash 43 | mkdir lst_directory 44 | cd lst_directory 45 | ``` 46 | 47 | 2. We recommend setting up a virtual environment for LST-AI: 48 | ```bash 49 | python3 -m venv /path/to/new/lst/virtual/environment 50 | ``` 51 | 52 | 3. Activate your new environment, e.g. `(lst_env)` 53 | ```bash 54 | source /path/to/new/lst/virtual/environment/bin/activate 55 | ``` 56 | 57 | 4. Install LST-AI (and yes, with `pip -e` option!): 58 | ```bash 59 | git clone https://github.com/CompImg/LST-AI/ 60 | cd LST-AI 61 | pip install -e . 62 | cd .. 63 | ``` 64 | 65 | 4. Install [HD-BET](https://github.com/MIC-DKFZ/HD-BET) 66 | ```bash 67 | git clone https://github.com/MIC-DKFZ/HD-BET 68 | cd HD-BET 69 | git checkout ae160681324d524db3578e4135bf781f8206e146 70 | pip install -e . 71 | cd .. 72 | ``` 73 | 74 | 5. Download or Compile and install greedy for your platform 75 | * 6.1 (Variant A): Download the pre-built greedy tool and place it into structure 76 | 1) Download the tool 77 | ```bash 78 | wget "https://github.com/CompImg/LST-AI/releases/download/v1.0.0/greedy" 79 | ``` 80 | 2) and ensure it is a findable path: 81 | ```bash 82 | chmod +x greedy 83 | mkdir ~/bin 84 | mv greedy ~/bin 85 | export PATH="$HOME/bin:$PATH" 86 | ``` 87 | Naturally, you can place the binary in ANY directory if you add it to your `.bashrc` and export the location to the `$PATH`. 88 | * 6.2 (Variant B): Compile, make, and install the greedy tool (you will need to install both, VTK and ITK) 89 | ``` 90 | apt-get update && apt-get install -y \ 91 | build-essential \ 92 | libpng-dev \ 93 | libtiff-dev \ 94 | uuid-dev \ 95 | make \ 96 | cmake \ 97 | g++ \ 98 | libgl1-mesa-dev 99 | 100 | wget https://github.com/InsightSoftwareConsortium/ITK/archive/refs/tags/v5.2.1.tar.gz 101 | tar -zxvf v5.2.1.tar.gz 102 | cd ITK-5.2.1 103 | mkdir build 104 | cd build 105 | cmake .. 106 | make -j$(nproc) 107 | make install 108 | 109 | wget https://www.vtk.org/files/release/9.1/VTK-9.1.0.tar.gz 110 | tar -xf VTK-9.1.0.tar.gz 111 | cmake .. 112 | make -j$(nproc) 113 | make install 114 | 115 | git clone https://github.com/pyushkevich/greedy greedy 116 | cmake ../greedy 117 | make -j$(nproc) 118 | make install 119 | ``` 120 | 121 | ### Usage of LST-AI 122 | 123 | Once installed, LST-AI can be used as a simple command line tool. LST-AI expects you to provide **zipped NIFTIs (*.nii.gz)** as input and assumes the input images **NOT** to be **skull-stripped**. If you already have skull-stripped images, **do not forget** to provide the **--skull-stripped** option, otherwise, the segmentation performance will be severely affected. 124 | 125 | LST-AI always requires you to provide a `--t1` T1w and `--flair` FLAIR image and to specify an output path for the segmentation results `--output`. If you would like to keep all processing files, for example, the segmentations and skull-stripped images in the MNI152 space, provide a directory via `--temp`. 126 | 127 | #### Example usage: 128 | ``` 129 | (lst_env) jqm@workstation: lst --t1 t1.nii.gz --flair flair.nii.gz --output /mnt/data/lst/results --temp /mnt/data/lst/processing 130 | ``` 131 | 132 | #### Modes 133 | 134 | We provide three different modes: 135 | 136 | 1. **Default Mode - Segmentation + Annotation**: In this mode, you only need to provide the T1w and FLAIR input images. LST-AI will automatically segment and annotate your lesions according to McDonald's criteria. 137 | 138 | 2. **Segmentation Only**: If you only care about the binary segmentation, and not about the annotation/class (perventricular, ...), this mode is for you. It will (only) save the binary segmentation mask. To execute it, provide the `--segment_only` flag to run it. 139 | 140 | 3. **Annotation Only**: If you already have a satisfactory binary segmentation mask for your T1w/FLAIR images, you can only use the annotation/region labeling function. Please provide your existing segmentation via `--existing_seg /path/to/binary/mask`, and provide the `--annotate_only` flag to run it. We assume that the lesion mask is provided in the FLAIR image space. 141 | 142 | #### Other (useful) settings 143 | 144 | - `--temp `: If you would like to access intermediate pipeline results such as the skull-stripped T1w, and FLAIR images in MNI152 space, please provide a temporary directory using this flag. Otherwise, we create a temporary directory on the fly and remove it once the pipeline has finished. 145 | - `--device`: Provide an integer value (e.g. `0`) for a GPU ID or `cpu` if you do not have access to a GPU. 146 | - `--stripped`: Bypass skull-stripping. Only use if your images are (actually) skull-stripped. We cannot handle a mixture (e.g. skull-stripped T1w, but non-skull-stripped FLAIR) as of now. 147 | - `--threshold`: Provide a value between `0` and `1` that defines the threshold which is applied to the lesion probability map generated by the ensemble network to create the binary lesion mask. The default setting is 0.5. 148 | - `--clipping`: This flag can be used to define lower and upper percentiles that are used as min & max for standardization of image intensities in pre-processing. Changing these parameters can have an effect on the sensitivity of the segmentation process (e.g., higher max can yield higher sensitivity). The default setting is `0.5 99.5`, which indicates that the min is defined as the 0.5 percentile and the max is defined as the 99.5 percentile. 149 | - `--probability_map`: Save the lesion probability maps of the ensemble network and of each individual 3D UNet model of the ensemble network. The `--temp` flag must be set, otherwise the files will be removed as they are stored in the temporary directory along with the intermediate pipeline results. 150 | 151 | 152 | ### Dockerfile and Dockerhub 153 | 154 | While the installation and usage require internet access to install python packages and to download the weights and atlas, we understand that some researchers prefer to use lst-ai offline. Thus, we have decided to provide lst-ai as a CPU-/GPU-enabled docker container, which can be compiled using our scripts (for instructions please check the docker directory). If, instead of building the docker yourself, you would just rather use it, you can pull it from dockerhub instead. 155 | 156 | While we used to maintain jqmcginnis/lst-ai_cpu, we encourage everyone to use the unified jqmcginnis/lst-ai instead (CPU/GPU enabled), featuring the newest LST-AI version. 157 | You can pull it from dockerhub via executing: 158 | 159 | ```bash 160 | docker pull jqmcginnis/lst-ai:v1.2.0 161 | ``` 162 | 163 | ### Running the LST-AI Docker Container 164 | Once you have pulled (or built) your Docker image using the Dockerfile provided you can run the container using the `docker run` command. Here are the steps to bind mount your files and retrieve the results: 165 | 166 | #### Run the Docker Container with Bind Mounts 167 | The primary mechanism for sharing files between your host system and the Docker container is the `-v` or `--volume` flag, which specifies a bind mount. 168 | 169 | Here's a breakdown of how to use bind mounts: 170 | ```bash 171 | docker run -v [path_on_host]:[path_in_container] [image_name] 172 | ``` 173 | 174 | Given our provided GPU Dockerfile command, the run command might look something like this: 175 | 176 | ```bash 177 | docker run -v /home/ginnis/lst_in:/custom_apps/lst_input -v /home/ginnis/lst_out/:/custom_apps/lst_output -v /home/ginnis/lst_temp/:/custom_apps/lst_temp lst:latest --t1 /custom_apps/lst_input/t1.nii.gz --flair /custom_apps/lst_input/flair3d.nii.gz --output /custom_apps/lst_output --temp /custom_apps/lst_temp 178 | ``` 179 | 180 | __Note__: Ensure your paths are absolute, as Docker requires absolute paths for bind mounts. Since you've bind-mounted your output directory to `/home/ginnis/lst_out/` on your host, the results from the Docker container will be written directly to this directory. No additional steps are needed to retrieve the results, they will appear in this directory after the container has finished processing. 181 | 182 | #### Extending and modifying LST-AI for your custom code and pipeline 183 | 184 | We invite you to tailor LST-AI to your pipeline and application, please have a look at our [sources](LST-AI). 185 | 186 | ### BIDS Compliance with LST-AI 187 | 188 | To ensure maximum flexibility for our user base, LST-AI does not natively enforce BIDS-compliant file-naming conventions. This decision allows users to work seamlessly with both BIDS and non-BIDS datasets. 189 | 190 | However, for those who wish to utilize LST-AI within a BIDS-compliant workflow, we have provided an [example repository](https://github.com/twiltgen/LST-AI_BIDS) that demonstrates the integration of LST-AI with BIDS-compliant data. This example reflects the BIDS-compliant usage of LST-AI that we are currently using in our internal database. 191 | 192 | ### Citation 193 | 194 | Please consider citing [LST-AI](https://www.medrxiv.org/content/10.1101/2023.11.23.23298966) to support the development: 195 | ``` 196 | @article{wiltgen2024lst, 197 | title={LST-AI: A deep learning ensemble for accurate MS lesion segmentation}, 198 | author={Wiltgen, Tun and McGinnis, Julian and Schlaeger, Sarah and Kofler, Florian and Voon, CuiCi and Berthele, Achim and Bischl, Daria and Grundl, Lioba and Will, Nikolaus and Metz, Marie and others}, 199 | journal={NeuroImage: Clinical}, 200 | pages={103611}, 201 | year={2024}, 202 | publisher={Elsevier} 203 | } 204 | ``` 205 | 206 | Further, please also credit [greedy](https://greedy.readthedocs.io/en/latest/), and [HD-BET](https://github.com/MIC-DKFZ/HD-BET) used for preprocessing the image data. 207 | 208 | greedy 209 | ``` 210 | @article{yushkevich2016ic, 211 | title={IC-P-174: Fast Automatic Segmentation of Hippocampal Subfields and Medial Temporal Lobe Subregions In 3 Tesla and 7 Tesla T2-Weighted MRI}, 212 | author={Yushkevich, Paul A and Pluta, John and Wang, Hongzhi and Wisse, Laura EM and Das, Sandhitsu and Wolk, David}, 213 | journal={Alzheimer's \& Dementia}, 214 | volume={12}, 215 | pages={P126--P127}, 216 | year={2016}, 217 | publisher={Wiley Online Library} 218 | } 219 | ``` 220 | 221 | HD-BET: 222 | ``` 223 | @article{isensee2019automated, 224 | title={Automated brain extraction of multisequence MRI using artificial neural networks}, 225 | author={Isensee, Fabian and Schell, Marianne and Pflueger, Irada and Brugnara, Gianluca and Bonekamp, David and Neuberger, Ulf and Wick, Antje and Schlemmer, Heinz-Peter and Heiland, Sabine and Wick, Wolfgang and others}, 226 | journal={Human brain mapping}, 227 | volume={40}, 228 | number={17}, 229 | pages={4952--4964}, 230 | year={2019}, 231 | publisher={Wiley Online Library} 232 | } 233 | ``` 234 | -------------------------------------------------------------------------------- /docker/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM nvidia/cuda:12.3.1-runtime-ubuntu22.04 2 | 3 | # Prevents prompts from asking for user input during package installation 4 | ENV DEBIAN_FRONTEND=noninteractive 5 | 6 | # Update and install required packages 7 | RUN apt-get update && apt-get install -y \ 8 | git \ 9 | wget \ 10 | libeigen3-dev \ 11 | unzip 12 | 13 | # copied from https://stackoverflow.com/a/76170605/3485363 14 | RUN apt-get update && DEBIAN_FRONTEND=noninteractive \ 15 | apt-get install -y software-properties-common && \ 16 | add-apt-repository -y ppa:deadsnakes/ppa && \ 17 | apt-get install -y python3.10 curl && \ 18 | curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10 19 | 20 | RUN pip3 install --upgrade requests 21 | RUN ln -fs /usr/bin/python3.10 /usr/bin/python 22 | RUN python --version 23 | 24 | # # Setup LST-AI 25 | RUN mkdir -p /custom_apps/lst_directory 26 | 27 | # Install prerequisites 28 | # c.f. https://greedy.readthedocs.io/en/latest/install.html#compiling-from-source-code 29 | RUN apt-get update && \ 30 | apt-get install -y cmake g++ git 31 | 32 | # Install additional dependencies for VTK 33 | RUN apt-get install -y libgl1-mesa-dev libxt-dev 34 | 35 | # Install libpng 36 | RUN apt-get install -y libpng-dev 37 | 38 | # Build VTK 39 | # Download and unpack VTK 40 | WORKDIR /VTK 41 | RUN git clone https://gitlab.kitware.com/vtk/vtk.git 42 | WORKDIR /VTK/vtk 43 | RUN git checkout v9.1.0 44 | 45 | # Create and navigate to the build directory for VTK 46 | RUN mkdir VTK-build 47 | WORKDIR /VTK/vtk/VTK-build 48 | # ENV LD_LIBRARY_PATH=/VTK/vtk/VTK-build:$LD_LIBRARY_PATH 49 | 50 | # Run CMake to configure and build VTK 51 | RUN cmake -DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=OFF .. 52 | RUN make -j ${BUILD_JOBS} 53 | RUN make install 54 | 55 | # Build ITK 56 | # c.f. https://itk.org/Wiki/ITK/Getting_Started/Build/Linux 57 | # Clone the ITK repository 58 | RUN git clone https://github.com/InsightSoftwareConsortium/ITK.git /ITK 59 | WORKDIR /ITK 60 | 61 | # Checkout the specific version 62 | RUN git checkout v5.2.1 63 | 64 | # Create and navigate to the build directory 65 | RUN mkdir -p /ITK/build 66 | WORKDIR /ITK/build 67 | 68 | RUN apt-get install libexpat1-dev libgtest-dev libhdf5-dev libtiff-dev libvtkgdcm-dev -y 69 | 70 | # Run CMake to configure and build ITK 71 | RUN cmake -DModule_ITKPNG=ON \ 72 | -DBUILD_TESTING=OFF \ 73 | -DCMAKE_BUILD_TYPE=Release .. 74 | 75 | # run build process 76 | RUN make -j ${BUILD_JOBS} 77 | RUN make install 78 | 79 | # Clone the greedy repository 80 | RUN git clone https://github.com/pyushkevich/greedy /greedy 81 | WORKDIR /greedy 82 | RUN git checkout 1eafa4c6659b7a669fb299ce98d9531fc23e332a 83 | 84 | # Set the working directory to the build directory 85 | WORKDIR /greedy/build 86 | 87 | # Run ccmake from the build directory 88 | RUN cmake -DITK_DIR=/ITK/build \ 89 | -DVTK_DIR=/VTK/vtk/VTK-build \ 90 | -DCMAKE_BUILD_TYPE=Release \ 91 | -DBUILD_SHARED_LIBS=OFF \ 92 | .. 93 | RUN make -j ${BUILD_JOBS} 94 | RUN make install 95 | 96 | # Install HD-BET 97 | WORKDIR /custom_apps/lst_directory 98 | RUN git clone https://github.com/MIC-DKFZ/HD-BET 99 | WORKDIR /custom_apps/lst_directory/HD-BET 100 | RUN pip install -e . 101 | 102 | # Retrieve model weights for HD-BET 103 | WORKDIR /custom_apps/lst_directory/ 104 | RUN mkdir -p /root/hd-bet_params 105 | RUN wget -O /root/hd-bet_params/data.zip \ 106 | https://zenodo.org/api/records/2540695/files-archive 107 | WORKDIR /root/hd-bet_params/ 108 | RUN unzip data.zip && rm data.zip 109 | 110 | WORKDIR /custom_apps/lst_directory 111 | RUN apt-get install python3-dev -y 112 | RUN git clone https://github.com/CompImg/LST-AI/ 113 | 114 | WORKDIR /custom_apps/lst_directory/LST-AI 115 | RUN git pull origin main 116 | RUN git checkout v1.1.0 117 | 118 | # pip or pip3 depending on your system 119 | RUN pip install -e . 120 | 121 | # Retrieve model weights and files for LST-AI 122 | WORKDIR /custom_apps/lst_directory/ 123 | RUN wget -O /custom_apps/lst_directory/LST-AI/LST_AI/data.zip \ 124 | https://github.com/CompImg/LST-AI/releases/download/v1.1.0/lst_data.zip 125 | WORKDIR /custom_apps/lst_directory/LST-AI/LST_AI/ 126 | RUN unzip data.zip && rm data.zip 127 | 128 | # Make directories for easily mounting data 129 | # You may change these to your liking 130 | RUN mkdir -p /custom_apps/lst_input 131 | RUN mkdir -p /custom_apps/lst_output 132 | RUN mkdir -p /custom_apps/lst_temp 133 | 134 | # Entrypoint to run the python script when the container starts 135 | ENTRYPOINT [ "lst" ] -------------------------------------------------------------------------------- /docker/Readme.md: -------------------------------------------------------------------------------- 1 | #### Building the docker 2 | 3 | Info: We are happy to support both ARM64 and AMD64 platforms with the newest docker container. 4 | 5 | #### Guide on how to build the docker natively, and (tp push it to dockerhub) 6 | 7 | To build and push a Docker image for both linux/amd64 and linux/arm64/v8 platforms and then push it to Docker Hub under the name jqmcginnis/lst-ai, you can follow these steps: 8 | 9 | #### 1. Log in to dockerhub 10 | 11 | Open your terminal and log in to your Docker Hub account using the Docker CLI: 12 | 13 | ```bash 14 | docker login 15 | ``` 16 | Enter your Docker Hub username and password when prompted. 17 | 18 | #### 2. Enable Buildx (if not already enabled) 19 | 20 | Docker Buildx is an extended build feature that supports building multi-platform images. To ensure it is enabled, run: 21 | 22 | ```bash 23 | docker buildx create --use --name mybuilder 24 | ``` 25 | 26 | #### 3. Start a New Buildx Builder Instance 27 | 28 | This step ensures that the builder instance is started and uses the newly created builder: 29 | 30 | ```bash 31 | docker buildx use mybuilder 32 | docker buildx inspect --bootstrap 33 | ``` 34 | 35 | #### 4. Build and Push the Image 36 | 37 | Navigate to the directory where your Dockerfile is located, then build and push the image for both platforms. Replace path/to/dockerfile with the actual path to your Dockerfile if it's not in the current directory: 38 | 39 | ```bash 40 | docker buildx build --platform linux/amd64,linux/arm64/v8 -t jqmcginnis/lst-ai --push --build-arg BUILD_JOBS=8 . 41 | ``` 42 | This command will build the image for amd64 and arm64/v8 architectures and push it to Docker Hub under the repository jqmcginnis/lst-ai. It may take several hpurs (!). 43 | 44 | #### 5. Verify the Push 45 | 46 | Navigate to the directory where your Dockerfile is located, then build and push the image for both platforms. Replace path/to/dockerfile with the actual path to your Dockerfile if it's not in the current directory: 47 | 48 | ```bash 49 | docker buildx build --platform linux/amd64,linux/arm64/v8 -t jqmcginnis/lst-ai --push . 50 | ``` 51 | This command will build the image for amd64 and arm64/v8 architectures and push it to Docker Hub under the repository jqmcginnis/lst-ai. It may take several hours (!). 52 | -------------------------------------------------------------------------------- /figures/header.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CompImg/LST-AI/7d5ba1ca1be6f0a1ac7b72a14a247f999c2cc027/figures/header.png -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | setup( 4 | name='LST_AI', 5 | version='1.1.0', 6 | description='Lesion Segmentation Toolbox AI', 7 | url='https://github.com/CompImg/LST-AI', 8 | author='LST-AI Team', 9 | author_email=[ 10 | 'julian.mcginnis@tum.de', 11 | 'tun.wiltgen@tum.de', 12 | 'mark.muehlau@tum.de', 13 | 'benedict.wiestler@tum.de' 14 | ], 15 | keywords=['lesion_segmentation', 'ms', 'lst', 'ai'], 16 | python_requires='>=3.8,<3.10', 17 | install_requires=[ 18 | 'numpy<1.24.4', 19 | 'pillow', 20 | 'scipy>=1.9.0', 21 | 'scikit-image>=0.21.0', 22 | 'tensorflow>=2.13,<2.16', 23 | 'torch<=2.1.0', 24 | 'nibabel', 25 | 'requests' 26 | ], 27 | scripts=['LST_AI/lst'], 28 | license='MIT', 29 | packages=find_packages(include=['LST_AI']), 30 | classifiers=[ 31 | 'Intended Audience :: Science/Research', 32 | 'Programming Language :: Python', 33 | 'Topic :: Scientific/Engineering', 34 | 'Operating System :: Unix' 35 | ], 36 | ) 37 | --------------------------------------------------------------------------------