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

melage/{widgets/Synthstrip.py → plugins/warpseg/warpseg_main/voxelmorph/torch/Unet.py}

@@ -1,37 +1,84 @@
-"""
-Synthstrip model
-https://github.com/freesurfer/freesurfer/tree/dev/mri_synthstrip
-"""
-
 
+import torch.nn.functional as F
 import torch
 import torch.nn as nn
-import numpy as np
+import numpy as np  
+def default_unet_features():
+    nb_features = [
+        [16, 32, 32, 32],             # encoder
+        [32, 32, 32, 32, 32, 16, 16]  # decoder
+    ]
+    return nb_features
+class ConvBlock(nn.Module):
+    """
+    Specific convolutional block followed by leakyrelu for unet.
+    """
+
+    def __init__(self, ndims, in_channels, out_channels, stride=1):
+        super().__init__()
 
-class StripModel(nn.Module):
+        Conv = getattr(nn, 'Conv%dd' % ndims)
+        self.main = Conv(in_channels, out_channels, 3, stride, 1)
+        self.activation = nn.LeakyReLU(0.2)
+
+    def forward(self, x):
+        out = self.main(x)
+        out = self.activation(out)
+        return out
+
+class Unet(nn.Module):
     """
-    https://github.com/freesurfer/freesurfer/tree/dev/mri_synthstrip
+    A unet architecture. Layer features can be specified directly as a list of encoder and decoder
+    features or as a single integer along with a number of unet levels. The default network features
+    per layer (when no options are specified) are:
+
+        encoder: [16, 32, 32, 32]
+        decoder: [32, 32, 32, 32, 32, 16, 16]
     """
+
     def __init__(self,
-                 nb_features=16,
-                 nb_levels=7,
-                 feat_mult=2,
-                 max_features=64,
-                 nb_conv_per_level=2,
+                 inshape=None,
+                 infeats=None,
+                 nb_features=None,
+                 nb_levels=None,
                  max_pool=2,
-                 return_mask=False):
+                 feat_mult=1,
+                 nb_conv_per_level=1,
+                 half_res=False):
+        """
+        Parameters:
+            inshape: Input shape. e.g. (192, 192, 192)
+            infeats: Number of input features.
+            nb_features: Unet convolutional features. Can be specified via a list of lists with
+                the form [[encoder feats], [decoder feats]], or as a single integer.
+                If None (default), the unet features are defined by the default config described in
+                the class documentation.
+            nb_levels: Number of levels in unet. Only used when nb_features is an integer.
+                Default is None.
+            feat_mult: Per-level feature multiplier. Only used when nb_features is an integer.
+                Default is 1.
+            nb_conv_per_level: Number of convolutions per unet level. Default is 1.
+            half_res: Skip the last decoder upsampling. Default is False.
+        """
 
         super().__init__()
 
-        # dimensionality
-        ndims = 3
+        # ensure correct dimensionality
+        ndims = len(inshape)
+        assert ndims in [1, 2, 3], 'ndims should be one of 1, 2, or 3. found: %d' % ndims
+
+        # cache some parameters
+        self.half_res = half_res
+
+        # default encoder and decoder layer features if nothing provided
+        if nb_features is None:
+            nb_features = default_unet_features()
 
         # build feature list automatically
         if isinstance(nb_features, int):
             if nb_levels is None:
                 raise ValueError('must provide unet nb_levels if nb_features is an integer')
             feats = np.round(nb_features * feat_mult ** np.arange(nb_levels)).astype(int)
-            feats = np.clip(feats, 1, max_features)
             nb_features = [
                 np.repeat(feats[:-1], nb_conv_per_level),
                 np.repeat(np.flip(feats), nb_conv_per_level)
@@ -55,7 +102,7 @@ class StripModel(nn.Module):
         self.upsampling = [nn.Upsample(scale_factor=s, mode='nearest') for s in max_pool]
 
         # configure encoder (down-sampling path)
-        prev_nf = 1
+        prev_nf = infeats
         encoder_nfs = [prev_nf]
         self.encoder = nn.ModuleList()
         for level in range(self.nb_levels - 1):
@@ -77,7 +124,7 @@ class StripModel(nn.Module):
                 convs.append(ConvBlock(ndims, prev_nf, nf))
                 prev_nf = nf
             self.decoder.append(convs)
-            if level < (self.nb_levels - 1):
+            if not half_res or level < (self.nb_levels - 2):
                 prev_nf += encoder_nfs[level]
 
         # now we take care of any remaining convolutions
@@ -86,14 +133,10 @@ class StripModel(nn.Module):
             self.remaining.append(ConvBlock(ndims, prev_nf, nf))
             prev_nf = nf
 
-        # final convolutions
-        if return_mask:
-            self.remaining.append(ConvBlock(ndims, prev_nf, 2, activation=None))
-            self.remaining.append(nn.Softmax(dim=1))
-        else:
-            self.remaining.append(ConvBlock(ndims, prev_nf, 1, activation=None))
+        # cache final number of features
+        self.final_nf = prev_nf
 
-    def forward(self, x, a=None):
+    def forward(self, x):
 
         # encoder forward pass
         x_history = [x]
@@ -107,7 +150,7 @@ class StripModel(nn.Module):
         for level, convs in enumerate(self.decoder):
             for conv in convs:
                 x = conv(x)
-            if level < (self.nb_levels - 1):
+            if not self.half_res or level < (self.nb_levels - 2):
                 x = self.upsampling[level](x)
                 x = torch.cat([x, x_history.pop()], dim=1)
 
@@ -115,27 +158,4 @@ class StripModel(nn.Module):
         for conv in self.remaining:
             x = conv(x)
 
-        return x
-
-class ConvBlock(nn.Module):
-    """
-    Specific convolutional block followed by leakyrelu for unet.
-    """
-
-    def __init__(self, ndims, in_channels, out_channels, stride=1, activation='leaky'):
-        super().__init__()
-
-        Conv = getattr(nn, 'Conv%dd' % ndims)
-        self.conv = Conv(in_channels, out_channels, 3, stride, 1)
-        if activation == 'leaky':
-            self.activation = nn.LeakyReLU(0.2)
-        elif activation == None:
-            self.activation = None
-        else:
-            raise ValueError(f'Unknown activation: {activation}')
-
-    def forward(self, x):
-        out = self.conv(x)
-        if self.activation is not None:
-            out = self.activation(out)
-        return out
+        return x

melage/plugins/warpseg/warpseg_main/voxelmorph/torch/__init__.py

@@ -0,0 +1,4 @@
+from . import layers
+from . import networks
+from . import losses
+from . import utils

melage/plugins/warpseg/warpseg_main/voxelmorph/torch/layers.py

@@ -0,0 +1,97 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as nnf
+
+
+class SpatialTransformer(nn.Module):
+    """
+    N-D Spatial Transformer
+    """
+
+    def __init__(self, size, mode='bilinear'):
+        super().__init__()
+
+        self.mode = mode
+
+        # create sampling grid
+        vectors = [torch.arange(0, s) for s in size]
+        grids = torch.meshgrid(vectors)
+        grid = torch.stack(grids)
+        grid = torch.unsqueeze(grid, 0)
+        grid = grid.type(torch.FloatTensor)
+
+        # registering the grid as a buffer cleanly moves it to the GPU, but it also
+        # adds it to the state dict. this is annoying since everything in the state dict
+        # is included when saving weights to disk, so the model files are way bigger
+        # than they need to be. so far, there does not appear to be an elegant solution.
+        # see: https://discuss.pytorch.org/t/how-to-register-buffer-without-polluting-state-dict
+        self.register_buffer('grid', grid)
+
+    def forward(self, src, flow):
+        # new locations
+        new_locs = self.grid + flow
+        shape = flow.shape[2:]
+
+        # need to normalize grid values to [-1, 1] for resampler
+        for i in range(len(shape)):
+            new_locs[:, i, ...] = 2 * (new_locs[:, i, ...] / (shape[i] - 1) - 0.5)
+
+        # move channels dim to last position
+        # also not sure why, but the channels need to be reversed
+        if len(shape) == 2:
+            new_locs = new_locs.permute(0, 2, 3, 1)
+            new_locs = new_locs[..., [1, 0]]
+        elif len(shape) == 3:
+            new_locs = new_locs.permute(0, 2, 3, 4, 1)
+            new_locs = new_locs[..., [2, 1, 0]]
+
+        return nnf.grid_sample(src, new_locs, align_corners=True, mode=self.mode)
+
+
+class VecInt(nn.Module):
+    """
+    Integrates a vector field via scaling and squaring.
+    """
+
+    def __init__(self, inshape, nsteps):
+        super().__init__()
+
+        assert nsteps >= 0, 'nsteps should be >= 0, found: %d' % nsteps
+        self.nsteps = nsteps
+        self.scale = 1.0 / (2 ** self.nsteps)
+        self.transformer = SpatialTransformer(inshape)
+
+    def forward(self, vec):
+        vec = vec * self.scale
+        for _ in range(self.nsteps):
+            vec = vec + self.transformer(vec, vec)
+        return vec
+
+
+class ResizeTransform(nn.Module):
+    """
+    Resize a transform, which involves resizing the vector field *and* rescaling it.
+    """
+
+    def __init__(self, vel_resize, ndims):
+        super().__init__()
+        self.factor = 1.0 / vel_resize
+        self.mode = 'linear'
+        if ndims == 2:
+            self.mode = 'bi' + self.mode
+        elif ndims == 3:
+            self.mode = 'tri' + self.mode
+
+    def forward(self, x):
+        if self.factor < 1:
+            # resize first to save memory
+            x = nnf.interpolate(x, align_corners=True, scale_factor=self.factor, mode=self.mode)
+            x = self.factor * x
+
+        elif self.factor > 1:
+            # multiply first to save memory
+            x = self.factor * x
+            x = nnf.interpolate(x, align_corners=True, scale_factor=self.factor, mode=self.mode)
+
+        # don't do anything if resize is 1
+        return x