def test_init(self, input_size, cp):
model = layer.BSplines3DTransform(cp, input_size[1:-1])
if isinstance(cp, int):
cp = (cp, cp, cp)
assert model.cp_spacing == cp
def test_coefficients(self, input_size, cp):
filters = self.generate_filter_coefficients(cp_spacing=cp)
model = layer.BSplines3DTransform(cp, input_size[1:-1])
model.build(input_size)
assert np.allclose(filters, model.filter.numpy(), atol=1e-8)
def _resize_interpolate(self, field, control_points):
resize = layer.ResizeCPTransform(control_points)
field = resize(field)
interpolate = layer.BSplines3DTransform(control_points,
self.fixed_image_size)
field = interpolate(field)
return field
def test_build(self, input_size, cp):
model = layer.BSplines3DTransform(cp, input_size[1:-1])
model.build(input_size)
assert model.filter.shape == (
4 * cp[0],
4 * cp[1],
4 * cp[2],
3,
3,
)
def test_interpolation(self, input_size, cp):
model = layer.BSplines3DTransform(cp, input_size[1:-1])
model.build(input_size)
vol_shape = input_size[1:-1]
num_cp = ([input_size[0]] + [
int(np.ceil(isize / cpsize) + 3)
for isize, cpsize in zip(vol_shape, cp)
] + [input_size[-1]])
field = tf.random.normal(shape=num_cp, dtype=tf.float32)
ddf = model.call(field)
assert ddf.shape == input_size
def __init__(
self,
image_size: tuple,
out_channels: int,
num_channel_initial: int,
extract_levels: List[int],
out_kernel_initializer: str,
out_activation: str,
control_points: (tuple, None) = None,
**kwargs,
):
"""
Image is encoded gradually, i from level 0 to E,
then it is decoded gradually, j from level E to D.
Some of the decoded levels are used for generating extractions.
So, extract_levels are between [0, E] with E = max(extract_levels),
and D = min(extract_levels).
:param image_size: tuple, such as (dim1, dim2, dim3)
:param out_channels: int, number of channels for the extractions
:param num_channel_initial: int, number of initial channels.
:param extract_levels: list of int, number of extraction levels.
:param out_kernel_initializer: str, initializer to use for kernels.
:param out_activation: str, activation to use at end layer.
:param control_points: (tuple, None), specify the distance between control points (in voxels).
:param kwargs:
"""
super(LocalNet, self).__init__(**kwargs)
# save parameters
self._extract_levels = extract_levels
self._extract_max_level = max(self._extract_levels) # E
self._extract_min_level = min(self._extract_levels) # D
# init layer variables
num_channels = [
num_channel_initial * (2**level)
for level in range(self._extract_max_level + 1)
] # level 0 to E
self._downsample_blocks = [
layer.DownSampleResnetBlock(filters=num_channels[i],
kernel_size=7 if i == 0 else 3)
for i in range(self._extract_max_level)
] # level 0 to E-1
self._conv3d_block = layer.Conv3dBlock(
filters=num_channels[-1]) # level E
self._upsample_blocks = [
layer.LocalNetUpSampleResnetBlock(num_channels[level])
for level in range(self._extract_max_level -
1, self._extract_min_level - 1, -1)
] # level D to E-1
self._extract_layers = [
# if kernels are not initialized by zeros, with init NN, extract may be too large
layer.Conv3dWithResize(
output_shape=image_size,
filters=out_channels,
kernel_initializer=out_kernel_initializer,
activation=out_activation,
) for _ in self._extract_levels
]
self.resize = (layer.ResizeCPTransform(control_points)
if control_points is not None else False)
self.interpolate = (layer.BSplines3DTransform(control_points,
image_size)
if control_points is not None else False)
def __init__(
self,
image_size: tuple,
out_channels: int,
num_channel_initial: int,
depth: int,
out_kernel_initializer: str,
out_activation: str,
pooling: bool = True,
concat_skip: bool = False,
control_points: (tuple, None) = None,
**kwargs,
):
"""
Initialise UNet.
:param image_size: tuple, (dim1, dim2, dim3), dims of input image.
:param out_channels: int, number of channels for the output
:param num_channel_initial: int, number of initial channels
:param depth: int, input is at level 0, bottom is at level depth
:param out_kernel_initializer: str, which kernel to use as initializer
:param out_activation: str, activation at last layer
:param pooling: Boolean, for downsampling, use non-parameterized
pooling if true, otherwise use conv3d
:param concat_skip: Boolean, when upsampling, concatenate skipped
tensor if true, otherwise use addition
:param control_points: (tuple, None), specify the distance between control points (in voxels).
:param kwargs:
"""
super(UNet, self).__init__(**kwargs)
# init layer variables
num_channels = [num_channel_initial * (2 ** d) for d in range(depth + 1)]
self._num_channel_initial = num_channel_initial
self._depth = depth
self._downsample_blocks = [
layer.DownSampleResnetBlock(filters=num_channels[d], pooling=pooling)
for d in range(depth)
]
self._bottom_conv3d = layer.Conv3dBlock(filters=num_channels[depth])
self._bottom_res3d = layer.Residual3dBlock(filters=num_channels[depth])
self._upsample_blocks = [
layer.UpSampleResnetBlock(filters=num_channels[d], concat=concat_skip)
for d in range(depth)
]
self._output_conv3d = layer.Conv3dWithResize(
output_shape=image_size,
filters=out_channels,
kernel_initializer=out_kernel_initializer,
activation=out_activation,
)
self.resize = (
layer.ResizeCPTransform(control_points)
if control_points is not None
else False
)
self.interpolate = (
layer.BSplines3DTransform(control_points, image_size)
if control_points is not None
else False
)
