melage 0.0.65__py3-none-any.whl

melage/widgets/DeepLModels/NPP/models/utils.py

@@ -0,0 +1,303 @@
+# Copyright (c) 2021-2022, InterDigital Communications, Inc
+# All rights reserved.
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted (subject to the limitations in the disclaimer
+# below) provided that the following conditions are met:
+
+# * Redistributions of source code must retain the above copyright notice,
+#   this list of conditions and the following disclaimer.
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+# * Neither the name of InterDigital Communications, Inc nor the names of its
+#   contributors may be used to endorse or promote products derived from this
+#   software without specific prior written permission.
+
+# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
+# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
+# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT
+# NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
+# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
+# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
+# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+def find_named_module(module, query):
+    """Helper function to find a named module. Returns a `nn.Module` or `None`
+
+    Args:
+        module (nn.Module): the root module
+        query (str): the module name to find
+
+    Returns:
+        nn.Module or None
+    """
+
+    return next((m for n, m in module.named_modules() if n == query), None)
+
+
+def find_named_buffer(module, query):
+    """Helper function to find a named buffer. Returns a `torch.Tensor` or `None`
+
+    Args:
+        module (nn.Module): the root module
+        query (str): the buffer name to find
+
+    Returns:
+        torch.Tensor or None
+    """
+    return next((b for n, b in module.named_buffers() if n == query), None)
+
+
+def _update_registered_buffer(
+    module,
+    buffer_name,
+    state_dict_key,
+    state_dict,
+    policy="resize_if_empty",
+    dtype=torch.int,
+):
+    new_size = state_dict[state_dict_key].size()
+    registered_buf = find_named_buffer(module, buffer_name)
+
+    if policy in ("resize_if_empty", "resize"):
+        if registered_buf is None:
+            raise RuntimeError(f'buffer "{buffer_name}" was not registered')
+
+        if policy == "resize" or registered_buf.numel() == 0:
+            registered_buf.resize_(new_size)
+
+    elif policy == "register":
+        if registered_buf is not None:
+            raise RuntimeError(f'buffer "{buffer_name}" was already registered')
+
+        module.register_buffer(buffer_name, torch.empty(new_size, dtype=dtype).fill_(0))
+
+    else:
+        raise ValueError(f'Invalid policy "{policy}"')
+
+
+def update_registered_buffers(
+    module,
+    module_name,
+    buffer_names,
+    state_dict,
+    policy="resize_if_empty",
+    dtype=torch.int,
+):
+    """Update the registered buffers in a module according to the tensors sized
+    in a state_dict.
+
+    (There's no way in torch to directly load a buffer with a dynamic size)
+
+    Args:
+        module (nn.Module): the module
+        module_name (str): module name in the state dict
+        buffer_names (list(str)): list of the buffer names to resize in the module
+        state_dict (dict): the state dict
+        policy (str): Update policy, choose from
+            ('resize_if_empty', 'resize', 'register')
+        dtype (dtype): Type of buffer to be registered (when policy is 'register')
+    """
+    valid_buffer_names = [n for n, _ in module.named_buffers()]
+    for buffer_name in buffer_names:
+        if buffer_name not in valid_buffer_names:
+            raise ValueError(f'Invalid buffer name "{buffer_name}"')
+
+    for buffer_name in buffer_names:
+        _update_registered_buffer(
+            module,
+            buffer_name,
+            f"{module_name}.{buffer_name}",
+            state_dict,
+            policy,
+            dtype,
+        )
+
+
+def conv(in_channels, out_channels, kernel_size=5, stride=2):
+    return nn.Conv2d(
+        in_channels,
+        out_channels,
+        kernel_size=kernel_size,
+        stride=stride,
+        padding=kernel_size // 2,
+    )
+
+
+def deconv(in_channels, out_channels, kernel_size=5, stride=2):
+    return nn.ConvTranspose2d(
+        in_channels,
+        out_channels,
+        kernel_size=kernel_size,
+        stride=stride,
+        output_padding=stride - 1,
+        padding=kernel_size // 2,
+    )
+
+
+def quantize_ste(x):
+    """Differentiable quantization via the Straight-Through-Estimator."""
+    # STE (straight-through estimator) trick: x_hard - x_soft.detach() + x_soft
+    return (torch.round(x) - x).detach() + x
+
+
+def gaussian_kernel1d(
+    kernel_size: int, sigma: float, device: torch.device, dtype: torch.dtype
+):
+    """1D Gaussian kernel."""
+    khalf = (kernel_size - 1) / 2.0
+    x = torch.linspace(-khalf, khalf, steps=kernel_size, dtype=dtype, device=device)
+    pdf = torch.exp(-0.5 * (x / sigma).pow(2))
+    return pdf / pdf.sum()
+
+
+def gaussian_kernel2d(
+    kernel_size: int, sigma: float, device: torch.device, dtype: torch.dtype
+):
+    """2D Gaussian kernel."""
+    kernel = gaussian_kernel1d(kernel_size, sigma, device, dtype)
+    return torch.mm(kernel[:, None], kernel[None, :])
+
+
+def gaussian_blur(x, kernel=None, kernel_size=None, sigma=None):
+    """Apply a 2D gaussian blur on a given image tensor."""
+    if kernel is None:
+        if kernel_size is None or sigma is None:
+            raise RuntimeError("Missing kernel_size or sigma parameters")
+        dtype = x.dtype if torch.is_floating_point(x) else torch.float32
+        device = x.device
+        kernel = gaussian_kernel2d(kernel_size, sigma, device, dtype)
+
+    padding = kernel.size(0) // 2
+    x = F.pad(x, (padding, padding, padding, padding), mode="replicate")
+    x = torch.nn.functional.conv2d(
+        x,
+        kernel.expand(x.size(1), 1, kernel.size(0), kernel.size(1)),
+        groups=x.size(1),
+    )
+    return x
+
+
+def meshgrid2d(N: int, C: int, H: int, W: int, device: torch.device):
+    """Create a 2D meshgrid for interpolation."""
+    theta = torch.eye(2, 3, device=device).unsqueeze(0).expand(N, 2, 3)
+    return F.affine_grid(theta, (N, C, H, W), align_corners=False)
+
+
+
+
+def getscale(data, dst_min, dst_max, f_low=0.0, f_high=0.999):
+    """
+    Function to get offset and scale of image intensities to robustly rescale to range dst_min..dst_max.
+    Equivalent to how mri_convert conforms images.
+    :param np.ndarray data: image data (intensity values)
+    :param float dst_min: future minimal intensity value
+    :param float dst_max: future maximal intensity value
+    :param f_low: robust cropping at low end (0.0 no cropping)
+    :param f_high: robust cropping at higher end (0.999 crop one thousandths of high intensity voxels)
+    :return: float src_min: (adjusted) offset
+    :return: float scale: scale factor
+    """
+    # get min and max from source
+    src_min = np.min(data)
+    src_max = np.max(data)
+
+    print("Input:    min: " + format(src_min) + "  max: " + format(src_max))
+
+    if f_low == 0.0 and f_high == 1.0:
+        return src_min, 1.0
+
+    # compute non-zeros and total vox num
+    nz = (np.abs(data) >= 1e-15).sum()
+    voxnum = data.shape[0] * data.shape[1] * data.shape[2]
+
+    # compute histogram
+    histosize = 1000
+    bin_size = (src_max - src_min) / histosize
+    hist, bin_edges = np.histogram(data, histosize)
+
+    # compute cummulative sum
+    cs = np.concatenate(([0], np.cumsum(hist)))
+
+    # get lower limit
+    nth = int(f_low * voxnum)
+    idx = np.where(cs < nth)
+
+    if len(idx[0]) > 0:
+        idx = idx[0][-1] + 1
+
+    else:
+        idx = 0
+
+    src_min = idx * bin_size + src_min
+
+    # print("bin min: "+format(idx)+"  nth: "+format(nth)+"  passed: "+format(cs[idx])+"\n")
+    # get upper limit
+    nth = voxnum - int((1.0 - f_high) * nz)
+    idx = np.where(cs >= nth)
+
+    if len(idx[0]) > 0:
+        idx = idx[0][0] - 2
+
+    else:
+        print('ERROR: rescale upper bound not found')
+
+    src_max = idx * bin_size + src_min
+    # print("bin max: "+format(idx)+"  nth: "+format(nth)+"  passed: "+format(voxnum-cs[idx])+"\n")
+
+    # scale
+    if src_min == src_max:
+        scale = 1.0
+
+    else:
+        scale = (dst_max - dst_min) / (src_max - src_min)
+
+    print("rescale:  min: " + format(src_min) + "  max: " + format(src_max) + "  scale: " + format(scale))
+
+    return src_min, scale
+
+
+
+def scalecrop(data, dst_min, dst_max, src_min, scale):
+    """
+    Function to crop the intensity ranges to specific min and max values
+    :param np.ndarray data: Image data (intensity values)
+    :param float dst_min: future minimal intensity value
+    :param float dst_max: future maximal intensity value
+    :param float src_min: minimal value to consider from source (crops below)
+    :param float scale: scale value by which source will be shifted
+    :return: np.ndarray data_new: scaled image data
+    """
+    data_new = dst_min + scale * (data - src_min)
+
+    # clip
+    data_new = np.clip(data_new, dst_min, dst_max)
+    print("Output:   min: " + format(data_new.min()) + "  max: " + format(data_new.max()))
+
+    return data_new
+
+
+def normalize(img):
+    src_min, scale = getscale(img.data, 0, 255)
+
+    if not img.data.dtype == np.dtype(np.uint8):
+        new_data = scalecrop(img.data, 0, 255, src_min, scale)
+    else:
+        new_data = img.data
+
+    new_img = img
+    new_img.data = new_data/255
+
+    return new_img
+

melage/widgets/DeepLModels/NPP/npp.py

@@ -0,0 +1,116 @@
+import os
+import sys
+import torch
+import torch
+import torch.nn as nn
+import numpy as np
+import argparse
+import surfa as sf
+import scipy.ndimage
+from models.model import UNet
+from models.utils import normalize
+
+description = ''' 
+Neural Pre-processing (NPP) converts Head MRI images
+to an intensity-normalized, skull-stripped brain in a standard coordi-
+nate space. If you use NPP in your analysis, please cite:
+'''
+
+# os.system(f'python npp.py -i {inputfname} -o {outputdir}')
+
+# parse command line
+parser = argparse.ArgumentParser(description=description)
+parser.add_argument('-i', '--image', metavar='file', required=True, help='Input image to pre-processing.')
+parser.add_argument('-o', '--out', metavar='file', help='Save stripped image to path.')
+parser.add_argument('-w', '--weight', metavar='float', help='Smoothness of intensity normalization mapping. The range of smoothness is [-3,2],'
+                                                            ' where a larger value implies a higher degree of smoothing',default =-1)
+parser.add_argument('-g', '--gpu', action='store_true', help='Use the GPU.')
+
+if len(sys.argv) == 1:
+    parser.print_help()
+    exit(1)
+
+args = parser.parse_args()
+
+# sanity check on the inputs
+if not args.out:
+    sf.system.fatal('Must provide at least --out output flags.')
+elif not os.path.exists(os.path.dirname(args.out)):
+    sf.system.fatal('Output directory does not exist.')
+
+# check args.weight is in the range and float
+if args.weight:
+    args.weight = float(args.weight)
+    if args.weight < -3 or args.weight > 2:
+        sf.system.fatal('The range of smoothness should within [-3,2], where a larger value implies a higher degree of smoothing')
+
+# necessary for speed gains (I think)
+torch.backends.cudnn.benchmark = True
+torch.backends.cudnn.deterministic = True
+
+# configure GPU device
+if args.gpu:
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+    device = torch.device('cuda')
+    device_name = 'GPU'
+else:
+    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+    device = torch.device('cpu')
+    device_name = 'CPU'
+
+# configure model
+print(f'Configuring model on the {device_name}')
+
+
+with torch.no_grad():
+    model = UNet()
+    model.to(device)
+    model.eval()
+
+
+version = '1'
+print(f'Running Neural Pre-processing model version {version}')
+cwd = os.getcwd()
+modelfile = os.path.join(cwd,'models/checkpoints', f'npp_v{version}.pth')
+checkpoint = torch.load(modelfile, map_location=device)
+
+model.load_state_dict(checkpoint)
+
+# load input volume
+image = sf.load_volume(args.image)
+print(f'Input image read from: {args.image}')
+
+# frame check
+if image.nframes > 1:
+    sf.system.fatal('Input image cannot have more than 1 frame')
+
+#i normalize image to [0, 255] and to [0, 1]
+image = normalize(image)
+
+# conform image and fit to shape with factors of 64
+conformed = image.conform(voxsize=1.0, dtype='float32',shape=(256,256,256), method='nearest', orientation='LIA')
+
+# predict the surface distance transform
+with torch.no_grad():
+    input_tensor = torch.from_numpy(conformed.data[np.newaxis, np.newaxis]).to(device)
+    output = model(input_tensor,args.weight)
+    mni_norm = output[0].cpu().numpy().squeeze().astype(np.int16)
+    norm = output[1].cpu().numpy().squeeze().astype(np.int16)
+    scalar_field = output[2].cpu().numpy().squeeze().astype(np.int16)
+
+# unconform the sdt and extract mask
+mni_norm = conformed.new(mni_norm)#.resample_like(image,method='nearest', fill=0)
+norm = conformed.new(norm)#.resample_like(image, method='nearest',fill=0)
+scalar_field = conformed.new(scalar_field)#.resample_like(image, method='nearest',fill=0)
+
+# write the masked output
+if args.out:
+    filename = os.path.basename(args.image)
+    filename = filename.split('.')[0]
+    mni_norm.save(os.path.join(args.out,filename+'_mni_norm.nii.gz'))
+    norm.save(os.path.join(args.out,filename+'_norm.nii.gz'))
+    scalar_field.save(os.path.join(args.out,filename+'_scalar_field.nii.gz'))
+
+    print(f'Results saved to: {args.out}')
+
+print('If you use Neural Pre-processing in your analysis, please cite:')

melage/widgets/DeepLModels/NPP/requirements.txt

@@ -0,0 +1,8 @@
+torchio==0.18.63
+einops==0.4.0
+nibabel==3.2.1
+numpy==1.24.2
+pytorch_lightning==1.5.10
+torch==1.10.2
+SimpleITK==2.2.1
+surfa==0.4.2

melage/widgets/DeepLModels/NPP/train/train.py

@@ -0,0 +1,116 @@
+import argparse
+import random
+import shutil
+import sys
+import os
+import torch
+from torch.utils.data import DataLoader
+from models.model import NPP
+from dataset.mri_dataset_affine import Generate_dataset
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
+from pytorch_lightning.callbacks import RichProgressBar
+import glob
+
+def save_checkpoint(state, is_best, filename="checkpoint_dae.pth.tar"):
+    torch.save(state, filename)
+    if is_best:
+        shutil.copyfile(filename, filename+'_best_loss')
+
+def parse_args(argv):
+    parser = argparse.ArgumentParser(description="Example training script.")
+    parser.add_argument(
+        "-d", "--dataset", type=str, required=False, help="Training dataset"
+    )
+    parser.add_argument(
+        "-e",
+        "--epochs",
+        default=60,
+        type=int,
+        help="Number of epochs (default: %(default)s)",
+    )
+    parser.add_argument(
+        "-lr",
+        "--learning-rate",
+        default=1e-4,
+        type=float,
+        help="Learning rate (default: %(default)s)",
+    )
+    parser.add_argument(
+        "-n",
+        "--num-workers",
+        type=int,
+        default=8,
+        help="Dataloaders threads (default: %(default)s)",
+    )
+
+    parser.add_argument("--checkpoint", type=str, help="Path to a checkpoint",default="checkpoint_dae_{quality}.pth.tar".format(quality = 24))
+
+    args = parser.parse_args(argv)
+    return args
+
+
+def main(argv):
+    args = parse_args(argv)
+
+    train_dataset,val_dataset = Generate_dataset()
+
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_size=2,
+        num_workers=args.num_workers,
+        shuffle=True,
+        pin_memory=( "cuda"),persistent_workers=True
+    )
+
+    val_dataloader = DataLoader(
+        val_dataset,
+        batch_size=2,
+        num_workers=2,
+        shuffle=False,
+        pin_memory=( "cuda"),persistent_workers=True
+    )
+
+    root_dir_path = os.path.join('checkpoint', "np_128_tv_hyper6_mean")
+
+    trainer = pl.Trainer(
+        default_root_dir=root_dir_path,
+        devices=1,
+        max_epochs=args.epochs,
+        precision=16, accelerator="gpu",
+        callbacks=[
+            ModelCheckpoint(mode="max",every_n_epochs = 1),
+            LearningRateMonitor("epoch"),RichProgressBar(),
+        ],
+        benchmark=True,
+    )
+    trainer.logger._log_graph = True  # If True, we plot the computation graph in tensorboard
+    trainer.logger._default_hp_metric = None  # Optional logging argument that we don't need
+
+    # Check whether pretrained model exists. If yes, load it and skip training
+    files = glob.glob('checkpoint/np_128_tv_hyper5_mean/lightning_logs/*/checkpoints/*')
+
+    sorted_by_mtime_descending = sorted(files, key=lambda t: -os.stat(t).st_mtime)
+    if len(sorted_by_mtime_descending)>0:
+        pretrained_filename = sorted_by_mtime_descending[0]
+    else:
+        pretrained_filename = ''
+
+    resume_path = pretrained_filename
+    if os.path.isfile(pretrained_filename):
+        print("Found pretrained model at %s, loading..." % pretrained_filename)
+        # Automatically loads the model with the saved hyperparameters
+        state_dict = torch.load(pretrained_filename)
+        model = NPP(args.learning_rate)
+        model.load_state_dict(state_dict['state_dict'],strict=False)
+        trainer.fit(model, train_dataloader, val_dataloader)
+    else:
+        pl.seed_everything(42)  # To be reproducable
+        model = NPP(args.learning_rate)
+        trainer.fit(model, train_dataloader, val_dataloader,ckpt_path= resume_path)
+    return model
+
+
+if __name__ == "__main__":
+    #[1,1e-1,1e-2,1e-3]
+    main(sys.argv[1:])