def Conv(*args, **kwargs):
ndims = settings.get_ndims()
if ndims == 2:
return nn.Conv2d(*args, **kwargs)
elif ndims == 3:
return nn.Conv3d(*args, **kwargs)
else:
raise Exception()
def BatchNorm(*args, **kwargs):
ndims = settings.get_ndims()
if ndims == 2:
return nn.BatchNorm2d(*args, **kwargs)
elif ndims == 3:
return nn.BatchNorm3d(*args, **kwargs)
else:
raise Exception()
def interpol_mode(mode):
"""
returns an interpolation mode for the current dimensioanlity.
"""
ndims = settings.get_ndims()
if mode in ["linear", "bilinear", "bicubic", "trilinear"]:
mode = ["linear", "bilinear", "trilinear"][ndims - 1]
return mode
def __init__(self, mode="bilinear"):
"""
Instantiates the grid sampler.
The grid sampler samples a grid of values at coordinates.
Parameters:
mode: interpolation mode
"""
super().__init__()
self.mode = mode
self.ndims = settings.get_ndims()
def volumetric_image_to_tensor(img, dtype=torch.float32):
"""
transforms a 3d np.array indexed as H x W x D x C to a torch.tensor indexed as C x H x W x D
"""
ndims = settings.get_ndims()
if len(img.shape) == ndims:
# add channel-dim
img = np.expand_dims(img, -1)
# permute channel to front
permuted = np.moveaxis(img, -1, 0)
return torch.as_tensor(np.array(permuted), dtype=dtype)
def forward(self, flow):
# create identity grid
size = flow.shape[2:]
vectors = [
torch.arange(0, s, dtype=flow.dtype, device=flow.device)
for s in size
]
grids = torch.meshgrid(vectors)
identity = torch.stack(grids) # z, y, x
identity = identity.expand(flow.shape[0], *[-1] *
(settings.get_ndims() + 1)) # add batch
return identity
def __init__(self, mode="bilinear"):
"""
Instantiates the spatial transformer.
A spatial transformer transforms a src image with a flow of displacement vectors.
Parameters:
mode: interpolation mode
"""
super().__init__()
self.identity = Identity()
self.grid_sampler = GridSampler(mode=mode)
self.ndims = settings.get_ndims()
def Dropout(*args, **kwargs):
"""
performs channel-whise dropout. As described in the paper Efficient Object Localization Using Convolutional
Networks , if adjacent pixels within feature maps are strongly correlated (as is normally the case in early
convolution layers) then i.i.d. dropout will not regularize the activations and will otherwise just result
in an effective learning rate decrease.
"""
ndims = settings.get_ndims()
if ndims == 2:
return nn.Dropout2d(*args, **kwargs)
elif ndims == 3:
return nn.Dropout3d(*args, **kwargs)
else:
raise Exception()
def __init__(
self,
degrees=(-180, 180),
translate=(-0.5, 0.5),
scale=(0.9, 1.1),
shear=(-0.03, 0.03),
flip=True,
):
super().__init__()
self.degrees = degrees
self.translate = translate
self.scale = scale
self.shear = shear
self.flip = flip
self.ndims = settings.get_ndims()
self.itenditity = tnn.Identity()
self.transform = tnn.AffineSpatialTransformer()
def Upsample(size=None,
scale_factor=None,
mode="nearest",
align_corners=False):
ndims = settings.get_ndims()
mode = interpol_mode(mode)
if ndims == 2:
return nn.Upsample(size=size,
scale_factor=scale_factor,
mode=mode,
align_corners=align_corners)
elif ndims == 3:
return nn.Upsample(size=size,
scale_factor=scale_factor,
mode=mode,
align_corners=align_corners)
else:
raise Exception()
def __init__(self, in_channels):
super().__init__()
"""
instantiates the flow prediction layer.
Parameters:
in_channels: input channels
"""
ndims = settings.get_ndims()
# configure cnn
self.cnn = nn.Sequential(
tnn.Conv(in_channels, in_channels, kernel_size=3, padding=1),
nn.LeakyReLU(0.2),
tnn.Conv(in_channels, in_channels, kernel_size=3, padding=1),
nn.LeakyReLU(0.2),
tnn.Conv(in_channels, ndims, kernel_size=3, padding=1),
)
# init final cnn layer with small weights and bias
self.cnn[-1].weight = nn.Parameter(
Normal(0, 1e-5).sample(self.cnn[-1].weight.shape)
)
self.cnn[-1].bias = nn.Parameter(torch.zeros(self.cnn[-1].bias.shape))
def randomize(self, size):
"""
randomizes the transformation
Parameters:
size: Size of the data, eg [128, 128, 128]
"""
self.do_augment = np.random.rand() < self.p
if not self.do_augment:
return
# dimensionality stuff
C = settings.get_ndims()
# scale down H,W,D,...
size_small = [1, C] + [s // self.r for s in size]
# create some noise at lower resolution
flow = torch.randn(size_small) * self.m
# upsample to full resolution
mode = tnn.interpol_mode("linear")
flow = F.interpolate(flow, size=size, mode=mode, align_corners=False)
# integrate for smoothness and invertability
self.pos_flow = self.integrate(flow)
self.neg_flow = self.integrate(-flow)
