# coding=utf-8

from typing import List, Optional, Tuple, Union

import numpy as np

import tensorflow as tf

import tensorflow.keras.layers as tfkl

from deepreg.model import layer_util

from deepreg.model.backbone.u_net import UNet

from deepreg.registry import REGISTRY

class AffineHead(tfkl.Layer):

def __init__(

self,

image_size: tuple,

name: str = "AffineHead",

):

"""

Init.

:param image_size: such as (dim1, dim2, dim3)

:param name: name of the layer

"""

super().__init__(name=name)

self.reference_grid = layer_util.get_reference_grid(image_size)

self.transform_initial = tf.constant_initializer(

value=list(np.eye(4, 3).reshape((-1)))

)

self._flatten = tfkl.Flatten()

self._dense = tfkl.Dense(units=12, bias_initializer=self.transform_initial)

def call(

self, inputs: Union[tf.Tensor, List], **kwargs

) -> Tuple[tf.Tensor, tf.Tensor]:

"""

:param inputs: a tensor or a list of tensor with length 1

:param kwargs: additional args

:return: ddf and theta

- ddf has shape (batch, dim1, dim2, dim3, 3)

- theta has shape (batch, 4, 3)

"""

if isinstance(inputs, list):

inputs = inputs[0]

theta = self._dense(self._flatten(inputs))

theta = tf.reshape(theta, shape=(-1, 4, 3))

# warp the reference grid with affine parameters to output a ddf

grid_warped = layer_util.warp_grid(self.reference_grid, theta)

ddf = grid_warped - self.reference_grid

return ddf, theta

def get_config(self):

"""Return the config dictionary for recreating this class."""

config = super().get_config()

config.update(image_size=self.reference_grid.shape[:3])

return config

@REGISTRY.register_backbone(name="global")

class GlobalNet(UNet):

"""

Build GlobalNet for image registration.

GlobalNet is a special UNet where the decoder for up-sampling is skipped.

The network's outputs come from the bottom layer from the encoder directly.

Reference:

- Hu, Yipeng, et al.

"Label-driven weakly-supervised learning

for multimodal deformable image registration,"

https://arxiv.org/abs/1711.01666

"""

def __init__(

self,

image_size: tuple,

num_channel_initial: int,

extract_levels: Optional[Tuple[int, ...]] = None,

depth: Optional[int] = None,

name: str = "GlobalNet",

**kwargs,

):

"""

Image is encoded gradually, i from level 0 to E.

Then, a densely-connected layer outputs an affine

transformation.

:param image_size: tuple, such as (dim1, dim2, dim3)

:param num_channel_initial: int, number of initial channels

:param extract_levels: list, which levels from net to extract, deprecated.

If depth is not given, depth = max(extract_levels) will be used.

:param depth: depth of the encoder. If given, extract_levels is not used.

:param name: name of the backbone.

:param kwargs: additional arguments.

"""

if depth is None:

if extract_levels is None:

raise ValueError(

"GlobalNet requires `depth` or `extract_levels` "

"to define the depth of encoder. "

"If `depth` is not given, "

"the maximum value of `extract_levels` will be used."

"However the argument `extract_levels` is deprecated "

"and will be removed in future release."

)

depth = max(extract_levels)

super().__init__(

image_size=image_size,

num_channel_initial=num_channel_initial,

depth=depth,

extract_levels=(depth,),

name=name,

**kwargs,

)

def build_output_block(

self,

image_size: Tuple[int, ...],

extract_levels: Tuple[int, ...],

out_channels: int,

out_kernel_initializer: str,

out_activation: str,

) -> Union[tf.keras.Model, tfkl.Layer]:

"""

Build a block for output.

The input to this block is a list of length 1.

The output has two tensors.

:param image_size: such as (dim1, dim2, dim3)

:param extract_levels: not used

:param out_channels: not used

:param out_kernel_initializer: not used

:param out_activation: not used

:return: a block consists of one or multiple layers

"""

return AffineHead(image_size=image_size)