melage 0.0.65__py3-none-any.whl → 1.0.0__py3-none-any.whl

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

@@ -1,272 +0,0 @@
-import nibabel
-import torch
-from torch import nn
-import torch.nn.functional as F
-from pytorch_lightning import LightningModule, Trainer
-from melage.widgets.DeepLModels.NPP.models.losses import SSIM3D
-from einops import rearrange
-import numpy as np
-
-
-class FeedForward(nn.Module):
-    def __init__(self, dim, hidden_dim, dropout = 0.):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(dim, hidden_dim),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(hidden_dim, dim),
-            nn.Dropout(dropout)
-        )
-    def forward(self, x):
-        return self.net(x)
-
-
-class Attention(nn.Module):
-    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
-        super().__init__()
-        inner_dim = dim_head *  heads
-        project_out = not (heads == 1 and dim_head == dim)
-
-        self.heads = heads
-        self.scale = dim_head ** -0.5
-
-        self.attend = nn.Softmax(dim = -1)
-        self.dropout = nn.Dropout(dropout)
-
-        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
-
-        self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, dim),
-            nn.Dropout(dropout)
-        ) if project_out else nn.Identity()
-
-    def forward(self, x):
-        qkv = self.to_qkv(x).chunk(3, dim = -1)
-        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
-
-        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
-
-        attn = self.attend(dots)
-        attn = self.dropout(attn)
-
-        out = torch.matmul(attn, v)
-        out = rearrange(out, 'b h n d -> b n (h d)')
-        return self.to_out(out)
-
-class Grad:
-    """
-    N-D gradient loss.
-    """
-
-    def __init__(self, penalty='l1', loss_mult=None):
-        self.penalty = penalty
-        self.loss_mult = loss_mult
-
-    def __call__(self, y_pred):
-        dy = torch.abs(y_pred[:, :, 1:, :, :] - y_pred[:, :, :-1, :, :])
-        dx = torch.abs(y_pred[:, :, :, 1:, :] - y_pred[:, :, :, :-1, :])
-        dz = torch.abs(y_pred[:, :, :, :, 1:] - y_pred[:, :, :, :, :-1])
-
-        if self.penalty == 'l2':
-            dy = dy * dy
-            dx = dx * dx
-            dz = dz * dz
-
-        grad = torch.mean(dx) + torch.mean(dy) + torch.mean(dz)
-
-        return grad
-
-
-class VGGBlock(nn.Module):
-    def __init__(self, in_channels, middle_channels, out_channels,stride = 1,relu='lrelu'):
-        super().__init__()
-        if relu=='lrelu':
-            self.relu = nn.LeakyReLU(inplace=True,negative_slope=0.01)
-        else:
-            self.relu = nn.ReLU(inplace=True)
-
-        self.conv1 = nn.Conv3d(in_channels, middle_channels, 3, padding=1,stride=stride)
-        self.bn1 = nn.InstanceNorm3d(middle_channels,affine=True)
-        self.conv2 = nn.Conv3d(middle_channels, out_channels, 3, padding=1)
-        self.bn2 = nn.InstanceNorm3d(out_channels,affine=True)
-
-    def forward(self, x):
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-        out = self.relu(out)
-
-        return out
-
-
-class UNet(nn.Module):
-    def __init__(self, input_channels=1, output_channels=1,):
-        super().__init__()
-        self._leaky_relu_alpha = 0.01
-        nb_filter = [16, 32, 64, 128, 256,512]
-        self.nb_filter = nb_filter
-        self.pool = nn.MaxPool3d(2)
-        self.up = nn.Upsample(scale_factor=(2,2,2), mode='trilinear', align_corners=True)
-
-        self.conv0_0 = VGGBlock(input_channels, nb_filter[0], nb_filter[0])
-        self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1], nb_filter[1],stride=2)
-        self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2], nb_filter[2],stride=2)
-        self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3], nb_filter[3],stride=2)
-        self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4], nb_filter[4],stride=2)
-        self.conv5_0 = VGGBlock(nb_filter[4], nb_filter[5], nb_filter[5],stride=2)
-
-        self.conv4_1 = VGGBlock(nb_filter[4]+nb_filter[5], nb_filter[4], nb_filter[4])
-        self.conv3_1 = VGGBlock(nb_filter[3]+nb_filter[4], nb_filter[3], nb_filter[3])
-        self.conv2_1 = VGGBlock(nb_filter[2]+nb_filter[3], nb_filter[2], nb_filter[2])
-        self.conv1_1 = VGGBlock(nb_filter[1]+nb_filter[2], nb_filter[1], nb_filter[1])
-        self.conv0_1 = VGGBlock(nb_filter[0]+nb_filter[1], nb_filter[0], nb_filter[0])
-
-        self.final = nn.Conv3d(nb_filter[0], output_channels, kernel_size=1)
-        self.LN1 = nn.LayerNorm(nb_filter[5])
-        self.LN2 = nn.LayerNorm(nb_filter[5])
-        self.attention1 = Attention(nb_filter[5],heads = 8,dim_head=64)
-        self.mlp1 = FeedForward(nb_filter[5],nb_filter[5])
-        self.LN3 = nn.LayerNorm(nb_filter[5])
-        self.LN4 = nn.LayerNorm(nb_filter[5])
-        self.attention2 = Attention(nb_filter[5],heads = 8,dim_head=64)
-        self.mlp2 = FeedForward(nb_filter[5],nb_filter[5])
-
-        self.head = nn.Sequential(
-            nn.Linear(nb_filter[5], nb_filter[5] // 2, bias=False),
-            nn.ReLU(inplace=True),
-            nn.Linear(nb_filter[5] // 2, 12, bias=False),
-            nn.Tanh()
-        )
-        self.head2 = nn.Sequential(
-            nn.Linear(1, nb_filter[5]),
-            nn.ReLU(inplace=True),
-            nn.Linear(nb_filter[5], nb_filter[5]*4),
-            nn.ReLU(inplace=True),
-            nn.Linear(nb_filter[5]*4, sum(np.array(nb_filter[:-1]))*2),
-        )
-    def forward(self, input,weight=None):
-
-        input_downsampled = torch.nn.functional.interpolate(input, size=[128, 128, 128],
-                                                mode='trilinear', align_corners=False)
-
-        x0_0 = self.conv0_0(input_downsampled)
-        x1_0 = self.conv1_0(x0_0)
-        x2_0 = self.conv2_0(x1_0)
-        x3_0 = self.conv3_0(x2_0)
-        x4_0 = self.conv4_0(x3_0)
-        x5_0 = self.conv5_0(x4_0)
-
-        identity = torch.eye(3,4).repeat(x5_0.shape[0],1,1).type_as(x5_0)
-        b,c,h,w,d = x5_0.shape
-        x5_0_faltten = rearrange(x5_0,'b c h w d-> b (h w d) c')
-        x5_0_faltten = self.attention1(self.LN1(x5_0_faltten)) + x5_0_faltten
-        x5_0_faltten = self.mlp1(self.LN2(x5_0_faltten)) + x5_0_faltten
-        x5_0_faltten = self.attention2(self.LN3(x5_0_faltten)) + x5_0_faltten
-        x5_0_faltten = self.mlp2(self.LN4(x5_0_faltten)) + x5_0_faltten
-
-        affine = self.head(x5_0_faltten.mean(dim=1)).reshape(-1, 3, 4) + identity
-        x5_0 = rearrange(x5_0_faltten,'b (h w d) c -> b c h w d',h=h,w=w,d=d)
-
-        x0_0_warp = torch.nn.functional.affine_grid(affine, input.size(), align_corners=False)
-        mod = self.head2(torch.Tensor([weight]).type_as(x0_0_warp))
-        mod = torch.split(mod,np.repeat(self.nb_filter[:-1],2).tolist(),0)
-
-        x4_1 = self.conv4_1(torch.cat([x4_0, self.up(x5_0)], 1))*mod[-1].reshape(1,-1,1,1,1)+ mod[-2].reshape(1,-1,1,1,1)
-        x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_1)], 1))*mod[-3].reshape(1,-1,1,1,1) + mod[-4].reshape(1,-1,1,1,1)
-        x2_1 = self.conv2_1(torch.cat([x2_0, self.up(x3_1)], 1))*mod[-5].reshape(1,-1,1,1,1)+ mod[-6].reshape(1,-1,1,1,1)
-        x1_1 = self.conv1_1(torch.cat([x1_0, self.up(x2_1)], 1))*mod[-7].reshape(1,-1,1,1,1)+ mod[-8].reshape(1,-1,1,1,1)
-        x0_1 = self.conv0_1(torch.cat([x0_0, self.up(x1_1)], 1))*mod[-9].reshape(1,-1,1,1,1)+ mod[-10].reshape(1,-1,1,1,1)
- 
-        output = self.final(x0_1)
-        output_upsampled = torch.nn.functional.interpolate(output, size=[256, 256, 256],mode='trilinear', align_corners=False)
-        norm = input*255 * output_upsampled
-        mni_norm = torch.nn.functional.grid_sample(norm, x0_0_warp, align_corners=False)
-
-        return mni_norm,norm,output_upsampled
-class NPP(LightningModule): 
-    def __init__(
-        self,
-        lr,
-    ):
-        super().__init__()
-        self.save_hyperparameters()
-
-        # networks
-        self.generator = UNet()
-
-        self.ssim7 = SSIM3D(window_size=5)
-        self.mse = nn.L1Loss()
-        self.grad = Grad('l1')
-        self.automatic_optimization = False
-
-
-    def forward(self, x,ind = None):
-        return self.generator(x,ind)
-
-    def training_step(self, batch, batch_idx):
-        imgs, gts,brain_mask, ind = batch
-        opt = self.optimizers()
-        sch = self.lr_schedulers()
-        # train generator
-        # generate images
-        weight = np.random.uniform(2,-4)
-        self.generated_imgs = self(imgs,weight)
-
-        # log sampled images
-        loss_ssim_7 = -self.ssim7(self.generated_imgs[1], gts)
-
-        grad_loss = self.grad(self.generated_imgs[0])
-
-        loss = loss_ssim_7 + grad_loss*10**weight
-        opt.zero_grad()
-
-        self.manual_backward(loss)
-        opt.step()
-
-        if self.trainer.is_last_batch:
-            sch.step()
-
-        self.log("tv_train" , grad_loss,prog_bar=True, sync_dist=True)
-        self.log("ssim_7_train" , loss_ssim_7,prog_bar=True, sync_dist=True)
-        return loss
-
-    def configure_optimizers(self):
-        lr = self.hparams.lr
-        my_list = []
-        sparse_params = list(filter(lambda kv: kv[0] in my_list, self.generator.named_parameters()))
-        sparse_params = [i[1] for i in sparse_params]
-        base_params = list(filter(lambda kv: kv[0] not in my_list, self.generator.named_parameters()))
-        base_params = [i[1] for i in base_params]
-
-        base_params = [{"params": base_params},]
-        sparse_params = [{"params": sparse_params},]
-
-        optimizer = torch.optim.Adam(base_params, lr=lr)
-
-        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[30], gamma=0.5)
-        return [optimizer], [lr_scheduler,]
-
-    def validation_step(self, batch, batch_idx):
-        imgs, gts,brain_mask,ind = batch
-
-        # train generator
-        self.generated_imgs = self(imgs,0.1)
-
-        grad_loss = self.grad(self.generated_imgs[0])
-        loss_ssim_7 = -self.ssim7(self.generated_imgs[1],gts)
-        loss = grad_loss + loss_ssim_7
-        self.log("tv_val", grad_loss,prog_bar=True, sync_dist=True )
-        self.log("ssim_7_val", loss_ssim_7,prog_bar=True, sync_dist=True )
-
-        return loss
-
-
-
-
-
-
-

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

@@ -1,303 +0,0 @@
-# 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

@@ -1,116 +0,0 @@
-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

@@ -1,8 +0,0 @@
-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