DeepRegNet /
DeepReg
| @@ 44-89 (lines=46) @@ | ||
| 41 | 'expand_ratio': 6, 'id_skip': True, 'strides': 1, 'se_ratio': 0.25} |
|
| 42 | ] |
|
| 43 | ||
| 44 | class AffineHead(tfkl.Layer): |
|
| 45 | def __init__( |
|
| 46 | self, |
|
| 47 | image_size: tuple, |
|
| 48 | name: str = "AffineHead", |
|
| 49 | ): |
|
| 50 | """ |
|
| 51 | Init. |
|
| 52 | ||
| 53 | :param image_size: such as (dim1, dim2, dim3) |
|
| 54 | :param name: name of the layer |
|
| 55 | """ |
|
| 56 | super().__init__(name=name) |
|
| 57 | self.reference_grid = layer_util.get_reference_grid(image_size) |
|
| 58 | self.transform_initial = tf.constant_initializer( |
|
| 59 | value=list(np.eye(4, 3).reshape((-1))) |
|
| 60 | ) |
|
| 61 | self._flatten = tfkl.Flatten() |
|
| 62 | self._dense = tfkl.Dense(units=12, bias_initializer=self.transform_initial) |
|
| 63 | ||
| 64 | def call( |
|
| 65 | self, inputs: Union[tf.Tensor, List], **kwargs |
|
| 66 | ) -> Tuple[tf.Tensor, tf.Tensor]: |
|
| 67 | """ |
|
| 68 | ||
| 69 | :param inputs: a tensor or a list of tensor with length 1 |
|
| 70 | :param kwargs: additional args |
|
| 71 | :return: ddf and theta |
|
| 72 | ||
| 73 | - ddf has shape (batch, dim1, dim2, dim3, 3) |
|
| 74 | - theta has shape (batch, 4, 3) |
|
| 75 | """ |
|
| 76 | if isinstance(inputs, list): |
|
| 77 | inputs = inputs[0] |
|
| 78 | theta = self._dense(self._flatten(inputs)) |
|
| 79 | theta = tf.reshape(theta, shape=(-1, 4, 3)) |
|
| 80 | # warp the reference grid with affine parameters to output a ddf |
|
| 81 | grid_warped = layer_util.warp_grid(self.reference_grid, theta) |
|
| 82 | ddf = grid_warped - self.reference_grid |
|
| 83 | return ddf, theta |
|
| 84 | ||
| 85 | def get_config(self): |
|
| 86 | """Return the config dictionary for recreating this class.""" |
|
| 87 | config = super().get_config() |
|
| 88 | config.update(image_size=self.reference_grid.shape[:3]) |
|
| 89 | return config |
|
| 90 | ||
| 91 | ||
| 92 | @REGISTRY.register_backbone(name="efficient_net") |
|
| @@ 14-59 (lines=46) @@ | ||
| 11 | from deepreg.registry import REGISTRY |
|
| 12 | ||
| 13 | ||
| 14 | class AffineHead(tfkl.Layer): |
|
| 15 | def __init__( |
|
| 16 | self, |
|
| 17 | image_size: tuple, |
|
| 18 | name: str = "AffineHead", |
|
| 19 | ): |
|
| 20 | """ |
|
| 21 | Init. |
|
| 22 | ||
| 23 | :param image_size: such as (dim1, dim2, dim3) |
|
| 24 | :param name: name of the layer |
|
| 25 | """ |
|
| 26 | super().__init__(name=name) |
|
| 27 | self.reference_grid = layer_util.get_reference_grid(image_size) |
|
| 28 | self.transform_initial = tf.constant_initializer( |
|
| 29 | value=list(np.eye(4, 3).reshape((-1))) |
|
| 30 | ) |
|
| 31 | self._flatten = tfkl.Flatten() |
|
| 32 | self._dense = tfkl.Dense(units=12, bias_initializer=self.transform_initial) |
|
| 33 | ||
| 34 | def call( |
|
| 35 | self, inputs: Union[tf.Tensor, List], **kwargs |
|
| 36 | ) -> Tuple[tf.Tensor, tf.Tensor]: |
|
| 37 | """ |
|
| 38 | ||
| 39 | :param inputs: a tensor or a list of tensor with length 1 |
|
| 40 | :param kwargs: additional args |
|
| 41 | :return: ddf and theta |
|
| 42 | ||
| 43 | - ddf has shape (batch, dim1, dim2, dim3, 3) |
|
| 44 | - theta has shape (batch, 4, 3) |
|
| 45 | """ |
|
| 46 | if isinstance(inputs, list): |
|
| 47 | inputs = inputs[0] |
|
| 48 | theta = self._dense(self._flatten(inputs)) |
|
| 49 | theta = tf.reshape(theta, shape=(-1, 4, 3)) |
|
| 50 | # warp the reference grid with affine parameters to output a ddf |
|
| 51 | grid_warped = layer_util.warp_grid(self.reference_grid, theta) |
|
| 52 | ddf = grid_warped - self.reference_grid |
|
| 53 | return ddf, theta |
|
| 54 | ||
| 55 | def get_config(self): |
|
| 56 | """Return the config dictionary for recreating this class.""" |
|
| 57 | config = super().get_config() |
|
| 58 | config.update(image_size=self.reference_grid.shape[:3]) |
|
| 59 | return config |
|
| 60 | ||
| 61 | ||
| 62 | @REGISTRY.register_backbone(name="global") |
|
