deepreg.model.network.RegistrationModel.__init__() - Code Metrics - DeepRegNet/DeepReg - Measure and Improve Code Quality continuously with Scrutinizer

deepreg.model.network.RegistrationModel.init() A
last analyzed 2022-05-18 21:56 UTC

↳ Parent: deepreg.model.network

Complexity

Conditions

Size

Total Lines	39
Code Lines	22

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
eloc	22
dl	0
loc	39
rs	9.352
c	0
b	0
f	0
cc	1
nop	8

How to fix Many Parameters

import os
from abc import abstractmethod
from copy import deepcopy
from typing import Dict, Optional, Tuple

import tensorflow as tf

from deepreg import log
from deepreg.loss.label import compute_centroid_distance
from deepreg.model import layer, layer_util
from deepreg.model.backbone import GlobalNet
from deepreg.registry import REGISTRY

logger = log.get(__name__)


def dict_without(d: dict, key) -> dict:
    """
    Return a copy of the given dict without a certain key.

    :param d: dict to be copied.
    :param key: key to be removed.
    :return: the copy without a key
    """
    copied = deepcopy(d)
    copied.pop(key)
    return copied


class RegistrationModel(tf.keras.Model):
    """Interface for registration model."""

    def __init__(
        self,
        moving_image_size: Tuple,
        fixed_image_size: Tuple,
        index_size: int,
        labeled: bool,
        batch_size: int,
        config: dict,
        name: str = "RegistrationModel",
    ):
        """
        Init.

        :param moving_image_size: (m_dim1, m_dim2, m_dim3)
        :param fixed_image_size: (f_dim1, f_dim2, f_dim3)
        :param index_size: number of indices for identify each sample
        :param labeled: if the data is labeled
        :param batch_size: total number of samples consumed per step, over all devices.
            When using multiple devices, TensorFlow automatically split the tensors.
            Therefore, input shapes should be defined over batch_size.
        :param config: config for method, backbone, and loss.
        :param name: name of the model
        """
        super().__init__(name=name)
        self.moving_image_size = moving_image_size
        self.fixed_image_size = fixed_image_size
        self.index_size = index_size
        self.labeled = labeled
        self.config = config
        self.batch_size = batch_size

        self._inputs = None  # save inputs of self._model as dict
        self._outputs = None  # save outputs of self._model as dict

        self.grid_ref = layer_util.get_reference_grid(grid_size=fixed_image_size)[
            None, ...
        ]
        self._model: tf.keras.Model = self.build_model()
        self.build_loss()

    def get_config(self) -> dict:
        """Return the config dictionary for recreating this class."""
        return dict(
            moving_image_size=self.moving_image_size,
            fixed_image_size=self.fixed_image_size,
            index_size=self.index_size,
            labeled=self.labeled,
            batch_size=self.batch_size,
            config=self.config,
            name=self.name,
        )

    @abstractmethod
    def build_model(self):
        """Build the model to be saved as self._model."""

    def build_inputs(self) -> Dict[str, tf.keras.layers.Input]:
        """
        Build input tensors.

        :return: dict of inputs.
        """
        # (batch, m_dim1, m_dim2, m_dim3)
        moving_image = tf.keras.Input(
            shape=self.moving_image_size,
            batch_size=self.batch_size,
            name="moving_image",
        )
        # (batch, f_dim1, f_dim2, f_dim3)
        fixed_image = tf.keras.Input(
            shape=self.fixed_image_size,
            batch_size=self.batch_size,
            name="fixed_image",
        )
        # (batch, index_size)
        indices = tf.keras.Input(
            shape=(self.index_size,),
            batch_size=self.batch_size,
            name="indices",
        )

        if not self.labeled:
            return dict(
                moving_image=moving_image, fixed_image=fixed_image, indices=indices
            )

        # (batch, m_dim1, m_dim2, m_dim3)
        moving_label = tf.keras.Input(
            shape=self.moving_image_size,
            batch_size=self.batch_size,
            name="moving_label",
        )
        # (batch, m_dim1, m_dim2, m_dim3)
        fixed_label = tf.keras.Input(
            shape=self.fixed_image_size,
            batch_size=self.batch_size,
            name="fixed_label",
        )
        return dict(
            moving_image=moving_image,
            fixed_image=fixed_image,
            moving_label=moving_label,
            fixed_label=fixed_label,
            indices=indices,
        )

    def concat_images(
        self,
        moving_image: tf.Tensor,
        fixed_image: tf.Tensor,
        moving_label: Optional[tf.Tensor] = None,
    ) -> tf.Tensor:
        """
        Adjust image shape and concatenate them together.

        :param moving_image: registration source
        :param fixed_image: registration target
        :param moving_label: optional, only used for conditional model.
        :return:
        """
        images = []

        resize_layer = layer.Resize3d(shape=self.fixed_image_size)

        # (batch, m_dim1, m_dim2, m_dim3, 1)
        moving_image = tf.expand_dims(moving_image, axis=4)
        moving_image = resize_layer(moving_image)
        images.append(moving_image)

        # (batch, m_dim1, m_dim2, m_dim3, 1)
        fixed_image = tf.expand_dims(fixed_image, axis=4)
        images.append(fixed_image)

        # (batch, m_dim1, m_dim2, m_dim3, 1)
        if moving_label is not None:
            moving_label = tf.expand_dims(moving_label, axis=4)
            moving_label = resize_layer(moving_label)
            images.append(moving_label)

        # (batch, f_dim1, f_dim2, f_dim3, 2 or 3)
        images = tf.concat(images, axis=4)
        return images

    def _build_loss(self, name: str, inputs_dict: dict):
        """
        Build and add one weighted loss together with the metrics.

        :param name: name of loss, image / label / regularization.
        :param inputs_dict: inputs for loss function
        """

        if name not in self.config["loss"]:
            # loss config is not defined
            logger.warning(
                f"The configuration for loss {name} is not defined. "
                f"Therefore it is not used."
            )
            return

        loss_configs = self.config["loss"][name]
        if not isinstance(loss_configs, list):
            loss_configs = [loss_configs]

        for loss_config in loss_configs:

            if "weight" not in loss_config:
                # default loss weight 1
                logger.warning(
                    f"The weight for loss {name} is not defined."
                    f"Default weight = 1.0 is used."
                )
                loss_config["weight"] = 1.0

            # build loss
            weight = loss_config["weight"]

            if weight == 0:
                logger.warning(
                    f"The weight for loss {name} is zero." f"Loss is not used."
                )
                return

            # do not perform reduction over batch axis for supporting multi-device
            # training, model.fit() will average over global batch size automatically
            loss_layer: tf.keras.layers.Layer = REGISTRY.build_loss(
                config=dict_without(d=loss_config, key="weight"),
                default_args={"reduction": tf.keras.losses.Reduction.NONE},
            )
            loss_value = loss_layer(**inputs_dict)
            weighted_loss = loss_value * weight

            # add loss
            self._model.add_loss(weighted_loss)

            # add metric
            self._model.add_metric(
                loss_value, name=f"loss/{name}_{loss_layer.name}", aggregation="mean"
            )
            self._model.add_metric(
                weighted_loss,
                name=f"loss/{name}_{loss_layer.name}_weighted",
                aggregation="mean",
            )

    @abstractmethod
    def build_loss(self):
        """Build losses according to configs."""

        # input metrics
        fixed_image = self._inputs["fixed_image"]
        moving_image = self._inputs["moving_image"]
        self.log_tensor_stats(tensor=moving_image, name="moving_image")
        self.log_tensor_stats(tensor=fixed_image, name="fixed_image")

        # image loss, conditional model does not have this
        if "pred_fixed_image" in self._outputs:
            pred_fixed_image = self._outputs["pred_fixed_image"]
            self._build_loss(
                name="image",
                inputs_dict=dict(y_true=fixed_image, y_pred=pred_fixed_image),
            )

        if self.labeled:
            # input metrics
            fixed_label = self._inputs["fixed_label"]
            moving_label = self._inputs["moving_label"]
            self.log_tensor_stats(tensor=moving_label, name="moving_label")
            self.log_tensor_stats(tensor=fixed_label, name="fixed_label")

            # label loss
            pred_fixed_label = self._outputs["pred_fixed_label"]
            self._build_loss(
                name="label",
                inputs_dict=dict(y_true=fixed_label, y_pred=pred_fixed_label),
            )

            # additional label metrics
            tre = compute_centroid_distance(
                y_true=fixed_label, y_pred=pred_fixed_label, grid=self.grid_ref
            )
            self._model.add_metric(tre, name="metric/TRE", aggregation="mean")

    def call(
        self, inputs: Dict[str, tf.Tensor], training=None, mask=None
    ) -> Dict[str, tf.Tensor]:
        """
        Call the self._model.

        :param inputs: a dict of tensors.
        :param training: training or not.
        :param mask: maks for inputs.
        :return:
        """
        return self._model(inputs, training=training, mask=mask)  # pragma: no cover

    @abstractmethod
    def postprocess(
        self,
        inputs: Dict[str, tf.Tensor],
        outputs: Dict[str, tf.Tensor],
    ) -> Tuple[tf.Tensor, Dict]:
        """
        Return a dict used for saving inputs and outputs.

        :param inputs: dict of model inputs
        :param outputs: dict of model outputs
        :return: tuple, indices and a dict.
            In the dict, each value is (tensor, normalize, on_label), where
            - normalize = True if the tensor need to be normalized to [0, 1]
            - on_label = True if the tensor depends on label
        """

    def plot_model(self, output_dir: str):
        """
        Save model structure in png.

        :param output_dir: path to the output dir.
        """
        self._model.summary(print_fn=logger.debug)
        try:
            tf.keras.utils.plot_model(
                self._model,
                to_file=os.path.join(output_dir, f"{self.name}.png"),
                dpi=96,
                show_shapes=True,
                show_layer_names=True,
                expand_nested=False,
            )
        except ImportError as err:  # pragma: no cover
            logger.error(
                "Failed to plot model structure. "
                "Please check if graphviz is installed. "
                "Error message is: %s.",
                err,
            )

    def log_tensor_stats(self, tensor: tf.Tensor, name: str):
        """
        Log statistics of a given tensor.

        :param tensor: tensor to monitor.
        :param name: name of the tensor.
        """
        flatten = tf.reshape(tensor, shape=(self.batch_size, -1))
        self._model.add_metric(
            tf.reduce_mean(flatten, axis=1),
            name=f"metric/{name}_mean",
            aggregation="mean",
        )
        self._model.add_metric(
            tf.reduce_min(flatten, axis=1),
            name=f"metric/{name}_min",
            aggregation="min",
        )
        self._model.add_metric(
            tf.reduce_max(flatten, axis=1),
            name=f"metric/{name}_max",
            aggregation="max",
        )


@REGISTRY.register_model(name="ddf")
class DDFModel(RegistrationModel):
    """
    A registration model predicts DDF.

    When using global net as backbone,
    the model predicts an affine transformation parameters,
    and a DDF is calculated based on that.
    """

    name = "DDFModel"

    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 build_model(self):
        """Build the model to be saved as self._model."""
        # build inputs
        self._inputs = self.build_inputs()
        moving_image = self._inputs["moving_image"]  # (batch, m_dim1, m_dim2, m_dim3)
        fixed_image = self._inputs["fixed_image"]  # (batch, f_dim1, f_dim2, f_dim3)

        # build ddf
        control_points = self.config["backbone"].pop("control_points", False)
        backbone_inputs = self.concat_images(moving_image, fixed_image)
        backbone = REGISTRY.build_backbone(
            config=self.config["backbone"],
            default_args=dict(
                image_size=self.fixed_image_size,
                out_channels=3,
                out_kernel_initializer="zeros",
                out_activation=None,
            ),
        )

        if isinstance(backbone, GlobalNet):
            # (f_dim1, f_dim2, f_dim3, 3), (4, 3)
            ddf, theta = backbone(inputs=backbone_inputs)
            self._outputs = dict(ddf=ddf, theta=theta)
        else:
            # (f_dim1, f_dim2, f_dim3, 3)
            ddf = backbone(inputs=backbone_inputs)
            ddf = (
                self._resize_interpolate(ddf, control_points) if control_points else ddf
            )
            self._outputs = dict(ddf=ddf)

        # build outputs
        warping = layer.Warping(fixed_image_size=self.fixed_image_size)
        # (f_dim1, f_dim2, f_dim3)
        pred_fixed_image = warping(inputs=[ddf, moving_image])
        self._outputs["pred_fixed_image"] = pred_fixed_image

        if not self.labeled:
            return tf.keras.Model(inputs=self._inputs, outputs=self._outputs)

        # (f_dim1, f_dim2, f_dim3)
        moving_label = self._inputs["moving_label"]
        pred_fixed_label = warping(inputs=[ddf, moving_label])

        self._outputs["pred_fixed_label"] = pred_fixed_label
        return tf.keras.Model(inputs=self._inputs, outputs=self._outputs)

    def build_loss(self):
        """Build losses according to configs."""
        super().build_loss()

        # ddf loss and metrics
        ddf = self._outputs["ddf"]
        self._build_loss(name="regularization", inputs_dict=dict(inputs=ddf))
        self.log_tensor_stats(tensor=ddf, name="ddf")

    def postprocess(
        self,
        inputs: Dict[str, tf.Tensor],
        outputs: Dict[str, tf.Tensor],
    ) -> Tuple[tf.Tensor, Dict]:
        """
        Return a dict used for saving inputs and outputs.

        :param inputs: dict of model inputs
        :param outputs: dict of model outputs
        :return: tuple, indices and a dict.
            In the dict, each value is (tensor, normalize, on_label), where
            - normalize = True if the tensor need to be normalized to [0, 1]
            - on_label = True if the tensor depends on label
        """
        indices = inputs["indices"]
        processed = dict(
            moving_image=(inputs["moving_image"], True, False),
            fixed_image=(inputs["fixed_image"], True, False),
            ddf=(outputs["ddf"], True, False),
            pred_fixed_image=(outputs["pred_fixed_image"], True, False),
        )

        # save theta for affine model
        if "theta" in outputs:
            processed["theta"] = (outputs["theta"], None, None)  # type: ignore

        if not self.labeled:
            return indices, processed

        processed = {
            **dict(
                moving_label=(inputs["moving_label"], False, True),
                fixed_label=(inputs["fixed_label"], False, True),
                pred_fixed_label=(outputs["pred_fixed_label"], False, True),
            ),
            **processed,
        }

        return indices, processed


@REGISTRY.register_model(name="dvf")
class DVFModel(DDFModel):
    """
    A registration model predicts DVF.

    DDF is calculated based on DVF.
    """

    name = "DVFModel"

    def build_model(self):
        """Build the model to be saved as self._model."""
        # build inputs
        self._inputs = self.build_inputs()
        moving_image = self._inputs["moving_image"]
        fixed_image = self._inputs["fixed_image"]
        control_points = self.config["backbone"].pop("control_points", False)

        # build ddf
        backbone_inputs = self.concat_images(moving_image, fixed_image)
        backbone = REGISTRY.build_backbone(
            config=self.config["backbone"],
            default_args=dict(
                image_size=self.fixed_image_size,
                out_channels=3,
                out_kernel_initializer="zeros",
                out_activation=None,
            ),
        )
        dvf = backbone(inputs=backbone_inputs)
        dvf = self._resize_interpolate(dvf, control_points) if control_points else dvf
        ddf = layer.IntDVF(fixed_image_size=self.fixed_image_size)(dvf)

        # build outputs
        self._warping = layer.Warping(fixed_image_size=self.fixed_image_size)
        # (f_dim1, f_dim2, f_dim3, 3)
        pred_fixed_image = self._warping(inputs=[ddf, moving_image])

        self._outputs = dict(dvf=dvf, ddf=ddf, pred_fixed_image=pred_fixed_image)

        if not self.labeled:
            return tf.keras.Model(inputs=self._inputs, outputs=self._outputs)

        # (f_dim1, f_dim2, f_dim3, 3)
        moving_label = self._inputs["moving_label"]
        pred_fixed_label = self._warping(inputs=[ddf, moving_label])

        self._outputs["pred_fixed_label"] = pred_fixed_label
        return tf.keras.Model(inputs=self._inputs, outputs=self._outputs)

    def build_loss(self):
        """Build losses according to configs."""
        super().build_loss()

        # dvf metrics
        dvf = self._outputs["dvf"]
        self.log_tensor_stats(tensor=dvf, name="dvf")

    def postprocess(
        self,
        inputs: Dict[str, tf.Tensor],
        outputs: Dict[str, tf.Tensor],
    ) -> Tuple[tf.Tensor, Dict]:
        """
        Return a dict used for saving inputs and outputs.

        :param inputs: dict of model inputs
        :param outputs: dict of model outputs
        :return: tuple, indices and a dict.
            In the dict, each value is (tensor, normalize, on_label), where
            - normalize = True if the tensor need to be normalized to [0, 1]
            - on_label = True if the tensor depends on label
        """
        indices, processed = super().postprocess(inputs=inputs, outputs=outputs)
        processed["dvf"] = (outputs["dvf"], True, False)
        return indices, processed


@REGISTRY.register_model(name="conditional")
class ConditionalModel(RegistrationModel):
    """
    A registration model predicts fixed image label without DDF or DVF.
    """

    name = "ConditionalModel"

    def build_model(self):
        """Build the model to be saved as self._model."""
        assert self.labeled

        # build inputs
        self._inputs = self.build_inputs()
        moving_image = self._inputs["moving_image"]
        fixed_image = self._inputs["fixed_image"]
        moving_label = self._inputs["moving_label"]

        # build ddf
        backbone_inputs = self.concat_images(moving_image, fixed_image, moving_label)
        backbone = REGISTRY.build_backbone(
            config=self.config["backbone"],
            default_args=dict(
                image_size=self.fixed_image_size,
                out_channels=1,
                out_kernel_initializer="glorot_uniform",
                out_activation="sigmoid",
            ),
        )
        # (batch, f_dim1, f_dim2, f_dim3)
        pred_fixed_label = backbone(inputs=backbone_inputs)
        pred_fixed_label = tf.squeeze(pred_fixed_label, axis=4)

        self._outputs = dict(pred_fixed_label=pred_fixed_label)
        return tf.keras.Model(inputs=self._inputs, outputs=self._outputs)

    def postprocess(
        self,
        inputs: Dict[str, tf.Tensor],
        outputs: Dict[str, tf.Tensor],
    ) -> Tuple[tf.Tensor, Dict]:
        """
        Return a dict used for saving inputs and outputs.

        :param inputs: dict of model inputs
        :param outputs: dict of model outputs
        :return: tuple, indices and a dict.
            In the dict, each value is (tensor, normalize, on_label), where
            - normalize = True if the tensor need to be normalized to [0, 1]
            - on_label = True if the tensor depends on label
        """
        indices = inputs["indices"]
        processed = dict(
            moving_image=(inputs["moving_image"], True, False),
            fixed_image=(inputs["fixed_image"], True, False),
            pred_fixed_label=(outputs["pred_fixed_label"], True, True),
            moving_label=(inputs["moving_label"], False, True),
            fixed_label=(inputs["fixed_label"], False, True),
        )

        return indices, processed


1			import os
2			from abc import abstractmethod
3			from copy import deepcopy
4			from typing import Dict, Optional, Tuple
5
6			import tensorflow as tf
7
8			from deepreg import log
9			from deepreg.loss.label import compute_centroid_distance
10			from deepreg.model import layer, layer_util
11			from deepreg.model.backbone import GlobalNet
12			from deepreg.registry import REGISTRY
13
14			logger = log.get(__name__)
15
16
17			def dict_without(d: dict, key) -> dict:
18			"""
19			Return a copy of the given dict without a certain key.
20
21			:param d: dict to be copied.
22			:param key: key to be removed.
23			:return: the copy without a key
24			"""
25			copied = deepcopy(d)
26			copied.pop(key)
27			return copied
28
29
30			class RegistrationModel(tf.keras.Model):
31			"""Interface for registration model."""
32
33			def __init__(
34			self,
35			moving_image_size: Tuple,
36			fixed_image_size: Tuple,
37			index_size: int,
38			labeled: bool,
39			batch_size: int,
40			config: dict,
41			name: str = "RegistrationModel",
42			):
43			"""
44			Init.
45
46			:param moving_image_size: (m_dim1, m_dim2, m_dim3)
47			:param fixed_image_size: (f_dim1, f_dim2, f_dim3)
48			:param index_size: number of indices for identify each sample
49			:param labeled: if the data is labeled
50			:param batch_size: total number of samples consumed per step, over all devices.
51			When using multiple devices, TensorFlow automatically split the tensors.
52			Therefore, input shapes should be defined over batch_size.
53			:param config: config for method, backbone, and loss.
54			:param name: name of the model
55			"""
56			super().__init__(name=name)
57			self.moving_image_size = moving_image_size
58			self.fixed_image_size = fixed_image_size
59			self.index_size = index_size
60			self.labeled = labeled
61			self.config = config
62			self.batch_size = batch_size
63
64			self._inputs = None # save inputs of self._model as dict
65			self._outputs = None # save outputs of self._model as dict
66
67			self.grid_ref = layer_util.get_reference_grid(grid_size=fixed_image_size)[
68			None, ...
69			]
70			self._model: tf.keras.Model = self.build_model()
71			self.build_loss()
72
73			def get_config(self) -> dict:
74			"""Return the config dictionary for recreating this class."""
75			return dict(
76			moving_image_size=self.moving_image_size,
77			fixed_image_size=self.fixed_image_size,
78			index_size=self.index_size,
79			labeled=self.labeled,
80			batch_size=self.batch_size,
81			config=self.config,
82			name=self.name,
83			)
84
85			@abstractmethod
86			def build_model(self):
87			"""Build the model to be saved as self._model."""
88
89			def build_inputs(self) -> Dict[str, tf.keras.layers.Input]:
90			"""
91			Build input tensors.
92
93			:return: dict of inputs.
94			"""
95			# (batch, m_dim1, m_dim2, m_dim3)
96			moving_image = tf.keras.Input(
97			shape=self.moving_image_size,
98			batch_size=self.batch_size,
99			name="moving_image",
100			)
101			# (batch, f_dim1, f_dim2, f_dim3)
102			fixed_image = tf.keras.Input(
103			shape=self.fixed_image_size,
104			batch_size=self.batch_size,
105			name="fixed_image",
106			)
107			# (batch, index_size)
108			indices = tf.keras.Input(
109			shape=(self.index_size,),
110			batch_size=self.batch_size,
111			name="indices",
112			)
113
114			if not self.labeled:
115			return dict(
116			moving_image=moving_image, fixed_image=fixed_image, indices=indices
117			)
118
119			# (batch, m_dim1, m_dim2, m_dim3)
120			moving_label = tf.keras.Input(
121			shape=self.moving_image_size,
122			batch_size=self.batch_size,
123			name="moving_label",
124			)
125			# (batch, m_dim1, m_dim2, m_dim3)
126			fixed_label = tf.keras.Input(
127			shape=self.fixed_image_size,
128			batch_size=self.batch_size,
129			name="fixed_label",
130			)
131			return dict(
132			moving_image=moving_image,
133			fixed_image=fixed_image,
134			moving_label=moving_label,
135			fixed_label=fixed_label,
136			indices=indices,
137			)
138
139			def concat_images(
140			self,
141			moving_image: tf.Tensor,
142			fixed_image: tf.Tensor,
143			moving_label: Optional[tf.Tensor] = None,
144			) -> tf.Tensor:
145			"""
146			Adjust image shape and concatenate them together.
147
148			:param moving_image: registration source
149			:param fixed_image: registration target
150			:param moving_label: optional, only used for conditional model.
151			:return:
152			"""
153			images = []
154
155			resize_layer = layer.Resize3d(shape=self.fixed_image_size)
156
157			# (batch, m_dim1, m_dim2, m_dim3, 1)
158			moving_image = tf.expand_dims(moving_image, axis=4)
159			moving_image = resize_layer(moving_image)
160			images.append(moving_image)
161
162			# (batch, m_dim1, m_dim2, m_dim3, 1)
163			fixed_image = tf.expand_dims(fixed_image, axis=4)
164			images.append(fixed_image)
165
166			# (batch, m_dim1, m_dim2, m_dim3, 1)
167			if moving_label is not None:
168			moving_label = tf.expand_dims(moving_label, axis=4)
169			moving_label = resize_layer(moving_label)
170			images.append(moving_label)
171
172			# (batch, f_dim1, f_dim2, f_dim3, 2 or 3)
173			images = tf.concat(images, axis=4)
174			return images
175
176			def _build_loss(self, name: str, inputs_dict: dict):
177			"""
178			Build and add one weighted loss together with the metrics.
179
180			:param name: name of loss, image / label / regularization.
181			:param inputs_dict: inputs for loss function
182			"""
183
184			if name not in self.config["loss"]:
185			# loss config is not defined
186			logger.warning(
187			f"The configuration for loss {name} is not defined. "
188			f"Therefore it is not used."
189			)
190			return
191
192			loss_configs = self.config["loss"][name]
193			if not isinstance(loss_configs, list):
194			loss_configs = [loss_configs]
195
196			for loss_config in loss_configs:
197
198			if "weight" not in loss_config:
199			# default loss weight 1
200			logger.warning(
201			f"The weight for loss {name} is not defined."
202			f"Default weight = 1.0 is used."
203			)
204			loss_config["weight"] = 1.0
205
206			# build loss
207			weight = loss_config["weight"]
208
209			if weight == 0:
210			logger.warning(
211			f"The weight for loss {name} is zero." f"Loss is not used."
212			)
213			return
214
215			# do not perform reduction over batch axis for supporting multi-device
216			# training, model.fit() will average over global batch size automatically
217			loss_layer: tf.keras.layers.Layer = REGISTRY.build_loss(
218			config=dict_without(d=loss_config, key="weight"),
219			default_args={"reduction": tf.keras.losses.Reduction.NONE},
220			)
221			loss_value = loss_layer(**inputs_dict)
222			weighted_loss = loss_value * weight
223
224			# add loss
225			self._model.add_loss(weighted_loss)
226
227			# add metric
228			self._model.add_metric(
229			loss_value, name=f"loss/{name}_{loss_layer.name}", aggregation="mean"
230			)
231			self._model.add_metric(
232			weighted_loss,
233			name=f"loss/{name}_{loss_layer.name}_weighted",
234			aggregation="mean",
235			)
236
237			@abstractmethod
238			def build_loss(self):
239			"""Build losses according to configs."""
240
241			# input metrics
242			fixed_image = self._inputs["fixed_image"]
243			moving_image = self._inputs["moving_image"]
244			self.log_tensor_stats(tensor=moving_image, name="moving_image")
245			self.log_tensor_stats(tensor=fixed_image, name="fixed_image")
246
247			# image loss, conditional model does not have this
248			if "pred_fixed_image" in self._outputs:
249			pred_fixed_image = self._outputs["pred_fixed_image"]
250			self._build_loss(
251			name="image",
252			inputs_dict=dict(y_true=fixed_image, y_pred=pred_fixed_image),
253			)
254
255			if self.labeled:
256			# input metrics
257			fixed_label = self._inputs["fixed_label"]
258			moving_label = self._inputs["moving_label"]
259			self.log_tensor_stats(tensor=moving_label, name="moving_label")
260			self.log_tensor_stats(tensor=fixed_label, name="fixed_label")
261
262			# label loss
263			pred_fixed_label = self._outputs["pred_fixed_label"]
264			self._build_loss(
265			name="label",
266			inputs_dict=dict(y_true=fixed_label, y_pred=pred_fixed_label),
267			)
268
269			# additional label metrics
270			tre = compute_centroid_distance(
271			y_true=fixed_label, y_pred=pred_fixed_label, grid=self.grid_ref
272			)
273			self._model.add_metric(tre, name="metric/TRE", aggregation="mean")
274
275			def call(
276			self, inputs: Dict[str, tf.Tensor], training=None, mask=None
277			) -> Dict[str, tf.Tensor]:
278			"""
279			Call the self._model.
280
281			:param inputs: a dict of tensors.
282			:param training: training or not.
283			:param mask: maks for inputs.
284			:return:
285			"""
286			return self._model(inputs, training=training, mask=mask) # pragma: no cover
287
288			@abstractmethod
289			def postprocess(
290			self,
291			inputs: Dict[str, tf.Tensor],
292			outputs: Dict[str, tf.Tensor],
293			) -> Tuple[tf.Tensor, Dict]:
294			"""
295			Return a dict used for saving inputs and outputs.
296
297			:param inputs: dict of model inputs
298			:param outputs: dict of model outputs
299			:return: tuple, indices and a dict.
300			In the dict, each value is (tensor, normalize, on_label), where
301			- normalize = True if the tensor need to be normalized to [0, 1]
302			- on_label = True if the tensor depends on label
303			"""
304
305			def plot_model(self, output_dir: str):
306			"""
307			Save model structure in png.
308
309			:param output_dir: path to the output dir.
310			"""
311			self._model.summary(print_fn=logger.debug)
312			try:
313			tf.keras.utils.plot_model(
314			self._model,
315			to_file=os.path.join(output_dir, f"{self.name}.png"),
316			dpi=96,
317			show_shapes=True,
318			show_layer_names=True,
319			expand_nested=False,
320			)
321			except ImportError as err: # pragma: no cover
322			logger.error(
323			"Failed to plot model structure. "
324			"Please check if graphviz is installed. "
325			"Error message is: %s.",
326			err,
327			)
328
329			def log_tensor_stats(self, tensor: tf.Tensor, name: str):
330			"""
331			Log statistics of a given tensor.
332
333			:param tensor: tensor to monitor.
334			:param name: name of the tensor.
335			"""
336			flatten = tf.reshape(tensor, shape=(self.batch_size, -1))
337			self._model.add_metric(
338			tf.reduce_mean(flatten, axis=1),
339			name=f"metric/{name}_mean",
340			aggregation="mean",
341			)
342			self._model.add_metric(
343			tf.reduce_min(flatten, axis=1),
344			name=f"metric/{name}_min",
345			aggregation="min",
346			)
347			self._model.add_metric(
348			tf.reduce_max(flatten, axis=1),
349			name=f"metric/{name}_max",
350			aggregation="max",
351			)
352
353
354			@REGISTRY.register_model(name="ddf")
355			class DDFModel(RegistrationModel):
356			"""
357			A registration model predicts DDF.
358
359			When using global net as backbone,
360			the model predicts an affine transformation parameters,
361			and a DDF is calculated based on that.
362			"""
363
364			name = "DDFModel"
365
366			def _resize_interpolate(self, field, control_points):
367			resize = layer.ResizeCPTransform(control_points)
368			field = resize(field)
369
370			interpolate = layer.BSplines3DTransform(control_points, self.fixed_image_size)
371			field = interpolate(field)
372
373			return field
374
375			def build_model(self):
376			"""Build the model to be saved as self._model."""
377			# build inputs
378			self._inputs = self.build_inputs()
379			moving_image = self._inputs["moving_image"] # (batch, m_dim1, m_dim2, m_dim3)
380			fixed_image = self._inputs["fixed_image"] # (batch, f_dim1, f_dim2, f_dim3)
381
382			# build ddf
383			control_points = self.config["backbone"].pop("control_points", False)
384			backbone_inputs = self.concat_images(moving_image, fixed_image)
385			backbone = REGISTRY.build_backbone(
386			config=self.config["backbone"],
387			default_args=dict(
388			image_size=self.fixed_image_size,
389			out_channels=3,
390			out_kernel_initializer="zeros",
391			out_activation=None,
392			),
393			)
394
395			if isinstance(backbone, GlobalNet):
396			# (f_dim1, f_dim2, f_dim3, 3), (4, 3)
397			ddf, theta = backbone(inputs=backbone_inputs)
398			self._outputs = dict(ddf=ddf, theta=theta)
399			else:
400			# (f_dim1, f_dim2, f_dim3, 3)
401			ddf = backbone(inputs=backbone_inputs)
402			ddf = (
403			self._resize_interpolate(ddf, control_points) if control_points else ddf
404			)
405			self._outputs = dict(ddf=ddf)
406
407			# build outputs
408			warping = layer.Warping(fixed_image_size=self.fixed_image_size)
409			# (f_dim1, f_dim2, f_dim3)
410			pred_fixed_image = warping(inputs=[ddf, moving_image])
411			self._outputs["pred_fixed_image"] = pred_fixed_image
412
413			if not self.labeled:
414			return tf.keras.Model(inputs=self._inputs, outputs=self._outputs)
415
416			# (f_dim1, f_dim2, f_dim3)
417			moving_label = self._inputs["moving_label"]
418			pred_fixed_label = warping(inputs=[ddf, moving_label])
419
420			self._outputs["pred_fixed_label"] = pred_fixed_label
421			return tf.keras.Model(inputs=self._inputs, outputs=self._outputs)
422
423			def build_loss(self):
424			"""Build losses according to configs."""
425			super().build_loss()
426
427			# ddf loss and metrics
428			ddf = self._outputs["ddf"]
429			self._build_loss(name="regularization", inputs_dict=dict(inputs=ddf))
430			self.log_tensor_stats(tensor=ddf, name="ddf")
431
432			def postprocess(
433			self,
434			inputs: Dict[str, tf.Tensor],
435			outputs: Dict[str, tf.Tensor],
436			) -> Tuple[tf.Tensor, Dict]:
437			"""
438			Return a dict used for saving inputs and outputs.
439
440			:param inputs: dict of model inputs
441			:param outputs: dict of model outputs
442			:return: tuple, indices and a dict.
443			In the dict, each value is (tensor, normalize, on_label), where
444			- normalize = True if the tensor need to be normalized to [0, 1]
445			- on_label = True if the tensor depends on label
446			"""
447			indices = inputs["indices"]
448			processed = dict(
449			moving_image=(inputs["moving_image"], True, False),
450			fixed_image=(inputs["fixed_image"], True, False),
451			ddf=(outputs["ddf"], True, False),
452			pred_fixed_image=(outputs["pred_fixed_image"], True, False),
453			)
454
455			# save theta for affine model
456			if "theta" in outputs:
457			processed["theta"] = (outputs["theta"], None, None) # type: ignore
458
459			if not self.labeled:
460			return indices, processed
461
462			processed = {
463			**dict(
464			moving_label=(inputs["moving_label"], False, True),
465			fixed_label=(inputs["fixed_label"], False, True),
466			pred_fixed_label=(outputs["pred_fixed_label"], False, True),
467			),
468			**processed,
469			}
470
471			return indices, processed
472
473
474			@REGISTRY.register_model(name="dvf")
475			class DVFModel(DDFModel):
476			"""
477			A registration model predicts DVF.
478
479			DDF is calculated based on DVF.
480			"""
481
482			name = "DVFModel"
483
484			def build_model(self):
485			"""Build the model to be saved as self._model."""
486			# build inputs
487			self._inputs = self.build_inputs()
488			moving_image = self._inputs["moving_image"]
489			fixed_image = self._inputs["fixed_image"]
490			control_points = self.config["backbone"].pop("control_points", False)
491
492			# build ddf
493			backbone_inputs = self.concat_images(moving_image, fixed_image)
494			backbone = REGISTRY.build_backbone(
495			config=self.config["backbone"],
496			default_args=dict(
497			image_size=self.fixed_image_size,
498			out_channels=3,
499			out_kernel_initializer="zeros",
500			out_activation=None,
501			),
502			)
503			dvf = backbone(inputs=backbone_inputs)
504			dvf = self._resize_interpolate(dvf, control_points) if control_points else dvf
505			ddf = layer.IntDVF(fixed_image_size=self.fixed_image_size)(dvf)
506
507			# build outputs
508			self._warping = layer.Warping(fixed_image_size=self.fixed_image_size)
509			# (f_dim1, f_dim2, f_dim3, 3)
510			pred_fixed_image = self._warping(inputs=[ddf, moving_image])
511
512			self._outputs = dict(dvf=dvf, ddf=ddf, pred_fixed_image=pred_fixed_image)
513
514			if not self.labeled:
515			return tf.keras.Model(inputs=self._inputs, outputs=self._outputs)
516
517			# (f_dim1, f_dim2, f_dim3, 3)
518			moving_label = self._inputs["moving_label"]
519			pred_fixed_label = self._warping(inputs=[ddf, moving_label])
520
521			self._outputs["pred_fixed_label"] = pred_fixed_label
522			return tf.keras.Model(inputs=self._inputs, outputs=self._outputs)
523
524			def build_loss(self):
525			"""Build losses according to configs."""
526			super().build_loss()
527
528			# dvf metrics
529			dvf = self._outputs["dvf"]
530			self.log_tensor_stats(tensor=dvf, name="dvf")
531
532			def postprocess(
533			self,
534			inputs: Dict[str, tf.Tensor],
535			outputs: Dict[str, tf.Tensor],
536			) -> Tuple[tf.Tensor, Dict]:
537			"""
538			Return a dict used for saving inputs and outputs.
539
540			:param inputs: dict of model inputs
541			:param outputs: dict of model outputs
542			:return: tuple, indices and a dict.
543			In the dict, each value is (tensor, normalize, on_label), where
544			- normalize = True if the tensor need to be normalized to [0, 1]
545			- on_label = True if the tensor depends on label
546			"""
547			indices, processed = super().postprocess(inputs=inputs, outputs=outputs)
548			processed["dvf"] = (outputs["dvf"], True, False)
549			return indices, processed
550
551
552			@REGISTRY.register_model(name="conditional")
553			class ConditionalModel(RegistrationModel):
554			"""
555			A registration model predicts fixed image label without DDF or DVF.
556			"""
557
558			name = "ConditionalModel"
559
560			def build_model(self):
561			"""Build the model to be saved as self._model."""
562			assert self.labeled
563
564			# build inputs
565			self._inputs = self.build_inputs()
566			moving_image = self._inputs["moving_image"]
567			fixed_image = self._inputs["fixed_image"]
568			moving_label = self._inputs["moving_label"]
569
570			# build ddf
571			backbone_inputs = self.concat_images(moving_image, fixed_image, moving_label)
572			backbone = REGISTRY.build_backbone(
573			config=self.config["backbone"],
574			default_args=dict(
575			image_size=self.fixed_image_size,
576			out_channels=1,
577			out_kernel_initializer="glorot_uniform",
578			out_activation="sigmoid",
579			),
580			)
581			# (batch, f_dim1, f_dim2, f_dim3)
582			pred_fixed_label = backbone(inputs=backbone_inputs)
583			pred_fixed_label = tf.squeeze(pred_fixed_label, axis=4)
584
585			self._outputs = dict(pred_fixed_label=pred_fixed_label)
586			return tf.keras.Model(inputs=self._inputs, outputs=self._outputs)
587
588			def postprocess(
589			self,
590			inputs: Dict[str, tf.Tensor],
591			outputs: Dict[str, tf.Tensor],
592			) -> Tuple[tf.Tensor, Dict]:
593			"""
594			Return a dict used for saving inputs and outputs.
595
596			:param inputs: dict of model inputs
597			:param outputs: dict of model outputs
598			:return: tuple, indices and a dict.
599			In the dict, each value is (tensor, normalize, on_label), where
600			- normalize = True if the tensor need to be normalized to [0, 1]
601			- on_label = True if the tensor depends on label
602			"""
603			indices = inputs["indices"]
604			processed = dict(
605			moving_image=(inputs["moving_image"], True, False),
606			fixed_image=(inputs["fixed_image"], True, False),
607			pred_fixed_label=(outputs["pred_fixed_label"], True, True),
608			moving_label=(inputs["moving_label"], False, True),
609			fixed_label=(inputs["fixed_label"], False, True),
610			)
611
612			return indices, processed
613

DeepRegNet / DeepReg

deepreg.model.network.RegistrationModel.__init__() A last analyzed 2022-05-18 21:56 UTC

Complexity

Size

Duplication

Importance

How to fix Many Parameters

Many Parameters

Duplication Side-by-Side

Filter issues like

deepreg.model.network.RegistrationModel.init() A
last analyzed 2022-05-18 21:56 UTC