CLI improvements

Author: ReubenDo

jstabl/jstabl_control

@@ -0,0 +1,216 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import argparse
+import os
+from tqdm import tqdm
+import SimpleITK as sitk
+from urllib.request import urlopen
+
+# pytorch 
+import torch
+import torch.nn 
+from torch.utils.data import DataLoader
+from torchvision.transforms import Compose
+
+# TorchIO
+import torchio
+from torchio import ImagesDataset, Image, Subject, DATA
+from torchio.transforms import (    
+        ZNormalization,
+        CenterCropOrPad,    
+        ToCanonical,    
+        Resample    
+    )   
+
+import jstabl
+from jstabl.networks.UNetModalityMeanTogether import Generic_UNet
+from jstabl.utilities.sampling import GridSampler, GridAggregator 
+
+   
+# Define training and patches sampling parameters   
+patch_size = (128,128,128)
+
+
+MODALITIES = ['T1']
+
+
+def multi_inference(paths_dict, 
+                    model,
+                    transformation,
+                    device,
+                    opt):
+    print("[INFO] Loading model.")
+    subjects_dataset_inf = ImagesDataset(paths_dict, transform=transformation)
+
+    border = (0,0,0)
+
+    print("[INFO] Starting Inference.")
+    for index, batch in enumerate(subjects_dataset_inf):
+        original_shape = batch['T1'][DATA].shape[1:]
+        reference = sitk.ReadImage(opt.t1)
+
+        new_shape = []
+        for i,dim in enumerate(original_shape):
+            new_dim = dim if dim>patch_size[i] else patch_size[i]
+            new_shape.append(new_dim)
+
+        batch_pad = CenterCropOrPad(tuple(new_shape))(batch)
+        affine_pad = batch_pad['T1']['affine']
+
+
+        data = batch_pad['T1'][DATA]
+    
+        sampler = GridSampler(data, opt.window_size, border)
+        aggregator = GridAggregator(data, border)
+        loader = DataLoader(sampler, batch_size=1)
+
+        with torch.no_grad():
+            for batch_elemt in tqdm(loader):
+                locations = batch_elemt['location']
+                input_tensor = {'T1':batch_elemt['image'][:,:1,...].to(device)}
+                #input_tensor['all'] = batch_elemt['image'].to(device)
+                labels = 0.0
+                for fold in range(1,4):
+                    path_model = os.path.join(list(jstabl.__path__)[0], f"./pretrained/glioma_{fold}.pth")
+                    model.load_state_dict(torch.load(path_model, map_location=device))
+                    logits,_ = model(input_tensor)
+                    labels+= torch.nn.Softmax(1)(logits)
+                labels = labels.argmax(dim=1, keepdim=True)
+                outputs = labels
+                aggregator.add_batch(outputs, locations)
+
+        output = aggregator.output_array.astype(float)
+        output = torchio.utils.nib_to_sitk(output, affine_pad)
+        output = sitk.Resample(
+            output,
+            reference,
+            sitk.Transform(),
+            sitk.sitkNearestNeighbor,
+        )
+        sitk.WriteImage(output, opt.res)
+
+
+def main():
+    opt = parsing_data()
+
+    if torch.cuda.is_available():
+        print("[INFO] GPU available.")  
+    else:
+        print("[INFO] GPU isn't available. Using CPU instead.")
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    print("[INFO] Reading data.")
+    assert os.path.exists(opt.t1), 'T1 scan not found'
+
+    if opt.res is None:
+        opt.res = opt.t1.replace('.nii', '_seg.nii')
+
+    output_folder = os.path.dirname(opt.res)
+    if not os.path.exists(output_folder) and len(output_folder)>0:
+        os.makedirs(output_folder)
+
+    if opt.preprocess:
+        try:
+            from MRIPreprocessor.mri_preprocessor import Preprocessor
+            output_folder = os.path.dirname(opt.res)
+            ppr = Preprocessor({
+                        'T1':opt.t1},
+                        output_folder=output_folder,
+                        reference='T1')
+            ppr.run_pipeline()
+            t1 = os.path.join(output_folder,'skullstripping', 'T1.nii.gz')
+            
+        except ImportError:
+            raise ImportError('Please install MRIPreprocessor. Run: ' '\n'
+                            '\t pip install  git+https://github.com/ReubenDo/MRIPreprocessor#egg=MRIPreprocessor')
+    else:
+        t1 = opt.t1   
+   
+    # filing paths
+    paths_dict = [Subject(
+        Image('T1', t1, torchio.INTENSITY))]
+    
+
+    transform_inference = (
+            ToCanonical(),
+            ZNormalization(),
+            Resample(1),      
+        )   
+    transform_inference = Compose(transform_inference) 
+     
+    # MODEL 
+    norm_op_kwargs = {'eps': 1e-5, 'affine': True}  
+    net_nonlin = torch.nn.LeakyReLU   
+    net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}   
+
+    print("[INFO] Building model.")  
+    model= Generic_UNet(input_modalities=['T1', 'all'],   
+                        base_num_features=32,   
+                        num_classes=10,     
+                        num_pool=4, 
+                        num_conv_per_stage=2,   
+                        feat_map_mul_on_downscale=2,    
+                        conv_op=torch.nn.Conv3d,    
+                        norm_op=torch.nn.InstanceNorm3d,    
+                        norm_op_kwargs=norm_op_kwargs,  
+                        nonlin=net_nonlin,  
+                        nonlin_kwargs=net_nonlin_kwargs,      
+                        convolutional_pooling=False,     
+                        convolutional_upsampling=False,  
+                        final_nonlin=lambda x: x,
+                        input_features={'T1':1, 'all':4})
+    
+    for fold in range(1,4):
+        path_model = os.path.join(list(jstabl.__path__)[0], f"./pretrained/glioma_{fold}.pth")
+        if not os.path.isfile(path_model):
+            url = f"https://zenodo.org/record/4040853/files/glioma_{fold}.pth?download=1"
+            print("Downloading", url, "...")
+            data = urlopen(url).read()
+            with open(path_model, 'wb') as f:
+                f.write(data)
+
+    
+    model.to(device)
+    model.eval()
+
+    multi_inference(paths_dict, model, transform_inference, device, opt)
+    print(f"[INFO] Inference done. Segmentation saved here: {opt.res}")
+    print("Have a good day!")
+
+
+def parsing_data():
+    parser = argparse.ArgumentParser(
+        description='Joint Tissue and Glioma segmentation designed for controls')
+
+
+    parser.add_argument('-t1',
+                        type=str,
+                        required=True,
+                        help='Filename of the T1 scan')
+
+    parser.add_argument('-res',
+                        type=str,
+                        default=None,
+                        help='Filename of the OUTPUT segmentation')
+
+    parser.add_argument('--preprocess',
+                        action='store_true', 
+                        help='Tag to use for preprocessing the data (co-registration + skull-stripping)')
+
+
+    parser.add_argument('--window_size',
+                        type=tuple,
+                        default=patch_size,
+                        help='Patch size for patch-based inference')
+
+    opt = parser.parse_args()
+
+    return opt
+
+
+if __name__ == '__main__':
+    main()
+
+
+