keras-hub-nightly 0.16.1.dev202410200345__py3-none-any.whl → 0.19.0.dev202412070351__py3-none-any.whl

keras_hub/src/models/retinanet/feature_pyramid.py

@@ -1,5 +1,9 @@
+import math
+
 import keras
 
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
 
 class FeaturePyramid(keras.layers.Layer):
     """A Feature Pyramid Network (FPN) layer.
@@ -37,14 +41,18 @@ class FeaturePyramid(keras.layers.Layer):
     Args:
         min_level: int. The minimum level of the feature pyramid.
         max_level: int. The maximum level of the feature pyramid.
+        use_p5: bool. If True, uses the output of the last layer (`P5` from
+            Feature Pyramid Network) as input for creating coarser convolution
+            layers (`P6`, `P7`).  If False, uses the direct input `P5`
+            for creating coarser convolution  layers.
         num_filters: int. The number of filters in each feature map.
         activation: string or `keras.activations`. The activation function
             to be used in network.
             Defaults to `"relu"`.
-        kernel_initializer: `str` or `keras.initializers` initializer.
+        kernel_initializer: `str` or `keras.initializers`.
             The kernel initializer for the convolution layers.
             Defaults to `"VarianceScaling"`.
-        bias_initializer: `str` or `keras.initializers` initializer.
+        bias_initializer: `str` or `keras.initializers`.
             The bias initializer for the convolution layers.
             Defaults to `"zeros"`.
         batch_norm_momentum: float.
@@ -53,10 +61,10 @@ class FeaturePyramid(keras.layers.Layer):
         batch_norm_epsilon: float.
             The epsilon for the batch normalization layers.
             Defaults to `0.001`.
-        kernel_regularizer: `str` or `keras.regularizers` regularizer.
+        kernel_regularizer: `str` or `keras.regularizers`.
             The kernel regularizer for the convolution layers.
             Defaults to `None`.
-        bias_regularizer: `str` or `keras.regularizers` regularizer.
+        bias_regularizer: `str` or `keras.regularizers`.
             The bias regularizer for the convolution layers.
             Defaults to `None`.
         use_batch_norm: bool. Whether to use batch normalization.
@@ -69,6 +77,7 @@ class FeaturePyramid(keras.layers.Layer):
         self,
         min_level,
         max_level,
+        use_p5,
         num_filters=256,
         activation="relu",
         kernel_initializer="VarianceScaling",
@@ -78,6 +87,7 @@ class FeaturePyramid(keras.layers.Layer):
         kernel_regularizer=None,
         bias_regularizer=None,
         use_batch_norm=False,
+        data_format=None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -89,6 +99,7 @@ class FeaturePyramid(keras.layers.Layer):
         self.min_level = min_level
         self.max_level = max_level
         self.num_filters = num_filters
+        self.use_p5 = use_p5
         self.activation = keras.activations.get(activation)
         self.kernel_initializer = keras.initializers.get(kernel_initializer)
         self.bias_initializer = keras.initializers.get(bias_initializer)
@@ -103,8 +114,8 @@ class FeaturePyramid(keras.layers.Layer):
             self.bias_regularizer = keras.regularizers.get(bias_regularizer)
         else:
             self.bias_regularizer = None
-        self.data_format = keras.backend.image_data_format()
-        self.batch_norm_axis = -1 if self.data_format == "channels_last" else 1
+        self.data_format = standardize_data_format(data_format)
+        self.batch_norm_axis = -1 if data_format == "channels_last" else 1
 
     def build(self, input_shapes):
         input_shapes = {
@@ -117,7 +128,6 @@ class FeaturePyramid(keras.layers.Layer):
         }
         input_levels = [int(level[1]) for level in input_shapes]
         backbone_max_level = min(max(input_levels), self.max_level)
-
         # Build lateral layers
         self.lateral_conv_layers = {}
         for i in range(self.min_level, backbone_max_level + 1):
@@ -134,7 +144,11 @@ class FeaturePyramid(keras.layers.Layer):
                 dtype=self.dtype_policy,
                 name=f"lateral_conv_{level}",
             )
-            self.lateral_conv_layers[level].build(input_shapes[level])
+            self.lateral_conv_layers[level].build(
+                (None, None, None, input_shapes[level][-1])
+                if self.data_format == "channels_last"
+                else (None, input_shapes[level][1], None, None)
+            )
 
         self.lateral_batch_norm_layers = {}
         if self.use_batch_norm:
@@ -149,9 +163,9 @@ class FeaturePyramid(keras.layers.Layer):
                     )
                 )
                 self.lateral_batch_norm_layers[level].build(
-                    (None, None, None, 256)
+                    (None, None, None, self.num_filters)
                     if self.data_format == "channels_last"
-                    else (None, 256, None, None)
+                    else (None, self.num_filters, None, None)
                 )
 
         # Build output layers
@@ -171,9 +185,9 @@ class FeaturePyramid(keras.layers.Layer):
                 name=f"output_conv_{level}",
             )
             self.output_conv_layers[level].build(
-                (None, None, None, 256)
+                (None, None, None, self.num_filters)
                 if self.data_format == "channels_last"
-                else (None, 256, None, None)
+                else (None, self.num_filters, None, None)
             )
 
         # Build coarser layers
@@ -192,11 +206,18 @@ class FeaturePyramid(keras.layers.Layer):
                 dtype=self.dtype_policy,
                 name=f"coarser_{level}",
             )
-            self.output_conv_layers[level].build(
-                (None, None, None, 256)
-                if self.data_format == "channels_last"
-                else (None, 256, None, None)
-            )
+            if i == backbone_max_level + 1 and self.use_p5:
+                self.output_conv_layers[level].build(
+                    (None, None, None, input_shapes[f"P{i-1}"][-1])
+                    if self.data_format == "channels_last"
+                    else (None, input_shapes[f"P{i-1}"][1], None, None)
+                )
+            else:
+                self.output_conv_layers[level].build(
+                    (None, None, None, self.num_filters)
+                    if self.data_format == "channels_last"
+                    else (None, self.num_filters, None, None)
+                )
 
         # Build batch norm layers
         self.output_batch_norms = {}
@@ -212,9 +233,9 @@ class FeaturePyramid(keras.layers.Layer):
                     )
                 )
                 self.output_batch_norms[level].build(
-                    (None, None, None, 256)
+                    (None, None, None, self.num_filters)
                     if self.data_format == "channels_last"
-                    else (None, 256, None, None)
+                    else (None, self.num_filters, None, None)
                 )
 
         # The same upsampling layer is used for all levels
@@ -273,7 +294,11 @@ class FeaturePyramid(keras.layers.Layer):
 
         for i in range(backbone_max_level + 1, self.max_level + 1):
             level = f"P{i}"
-            feats_in = output_features[f"P{i-1}"]
+            feats_in = (
+                inputs[f"P{i-1}"]
+                if i == backbone_max_level + 1 and self.use_p5
+                else output_features[f"P{i-1}"]
+            )
             if i > backbone_max_level + 1:
                 feats_in = self.activation(feats_in)
             output_features[level] = (
@@ -283,7 +308,10 @@ class FeaturePyramid(keras.layers.Layer):
                 if self.use_batch_norm
                 else self.output_conv_layers[level](feats_in)
             )
-
+        output_features = {
+            f"P{i}": output_features[f"P{i}"]
+            for i in range(self.min_level, self.max_level + 1)
+        }
         return output_features
 
     def get_config(self):
@@ -293,7 +321,9 @@ class FeaturePyramid(keras.layers.Layer):
                 "min_level": self.min_level,
                 "max_level": self.max_level,
                 "num_filters": self.num_filters,
+                "use_p5": self.use_p5,
                 "use_batch_norm": self.use_batch_norm,
+                "data_format": self.data_format,
                 "activation": keras.activations.serialize(self.activation),
                 "kernel_initializer": keras.initializers.serialize(
                     self.kernel_initializer
@@ -320,34 +350,51 @@ class FeaturePyramid(keras.layers.Layer):
 
     def compute_output_shape(self, input_shapes):
         output_shape = {}
-        print(input_shapes)
         input_levels = [int(level[1]) for level in input_shapes]
         backbone_max_level = min(max(input_levels), self.max_level)
 
         for i in range(self.min_level, backbone_max_level + 1):
             level = f"P{i}"
             if self.data_format == "channels_last":
-                output_shape[level] = input_shapes[level][:-1] + (256,)
+                output_shape[level] = input_shapes[level][:-1] + (
+                    self.num_filters,
+                )
             else:
                 output_shape[level] = (
                     input_shapes[level][0],
-                    256,
+                    self.num_filters,
                 ) + input_shapes[level][1:3]
 
         intermediate_shape = input_shapes[f"P{backbone_max_level}"]
         intermediate_shape = (
             (
                 intermediate_shape[0],
-                intermediate_shape[1] // 2,
-                intermediate_shape[2] // 2,
-                256,
+                (
+                    int(math.ceil(intermediate_shape[1] / 2))
+                    if intermediate_shape[1] is not None
+                    else None
+                ),
+                (
+                    int(math.ceil(intermediate_shape[1] / 2))
+                    if intermediate_shape[1] is not None
+                    else None
+                ),
+                self.num_filters,
             )
             if self.data_format == "channels_last"
             else (
                 intermediate_shape[0],
-                256,
-                intermediate_shape[1] // 2,
-                intermediate_shape[2] // 2,
+                self.num_filters,
+                (
+                    int(math.ceil(intermediate_shape[1] / 2))
+                    if intermediate_shape[1] is not None
+                    else None
+                ),
+                (
+                    int(math.ceil(intermediate_shape[1] / 2))
+                    if intermediate_shape[1] is not None
+                    else None
+                ),
             )
         )
 
@@ -357,16 +404,32 @@ class FeaturePyramid(keras.layers.Layer):
             intermediate_shape = (
                 (
                     intermediate_shape[0],
-                    intermediate_shape[1] // 2,
-                    intermediate_shape[2] // 2,
-                    256,
+                    (
+                        int(math.ceil(intermediate_shape[1] / 2))
+                        if intermediate_shape[1] is not None
+                        else None
+                    ),
+                    (
+                        int(math.ceil(intermediate_shape[1] / 2))
+                        if intermediate_shape[1] is not None
+                        else None
+                    ),
+                    self.num_filters,
                 )
                 if self.data_format == "channels_last"
                 else (
                     intermediate_shape[0],
-                    256,
-                    intermediate_shape[1] // 2,
-                    intermediate_shape[2] // 2,
+                    self.num_filters,
+                    (
+                        int(math.ceil(intermediate_shape[1] / 2))
+                        if intermediate_shape[1] is not None
+                        else None
+                    ),
+                    (
+                        int(math.ceil(intermediate_shape[1] / 2))
+                        if intermediate_shape[1] is not None
+                        else None
+                    ),
                 )
             )
 

keras_hub/src/models/retinanet/non_max_supression.py

@@ -3,6 +3,7 @@ import math
 import keras
 from keras import ops
 
+# TODO: https://github.com/keras-team/keras-hub/issues/1965
 from keras_hub.src.bounding_box import converters
 from keras_hub.src.bounding_box import utils
 from keras_hub.src.bounding_box import validate_format

keras_hub/src/models/retinanet/prediction_head.py

@@ -0,0 +1,192 @@
+import keras
+
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
+
+class PredictionHead(keras.layers.Layer):
+    """A head for classification or bounding box regression predictions.
+
+    Args:
+        output_filters: int. The umber of convolution filters in the final layer.
+            The number of output channels determines the prediction type:
+                - **Classification**:
+                    `output_filters = num_anchors * num_classes`
+                    Predicts class probabilities for each anchor.
+                - **Bounding Box Regression**:
+                    `output_filters = num_anchors * 4` Predicts bounding box
+                    offsets (x1, y1, x2, y2) for each anchor.
+        num_filters: int. The number of convolution filters to use in the base
+            layer.
+        num_conv_layers: int. The number of convolution layers before the final
+            layer.
+        use_prior_probability: bool. Set to True to use prior probability in the
+            bias initializer for the final convolution layer.
+            Defaults to `False`.
+        prior_probability: float. The prior probability value to use for
+            initializing the bias. Only used if `use_prior_probability` is
+            `True`. Defaults to `0.01`.
+        kernel_initializer: `str` or `keras.initializers`. The kernel
+            initializer for the convolution layers. Defaults to
+            `"random_normal"`.
+        bias_initializer: `str` or `keras.initializers`. The bias initializer
+            for the convolution layers. Defaults to `"zeros"`.
+        kernel_regularizer: `str` or `keras.regularizers`. The kernel
+            regularizer for the convolution layers. Defaults to `None`.
+        bias_regularizer: `str` or `keras.regularizers`. The bias regularizer
+            for the convolution layers. Defaults to `None`.
+        use_group_norm: bool. Whether to use Group Normalization after
+            the convolution layers. Defaults to `False`.
+
+    Returns:
+        A function representing either the classification
+            or the box regression head depending on `output_filters`.
+    """
+
+    def __init__(
+        self,
+        output_filters,
+        num_filters,
+        num_conv_layers,
+        use_prior_probability=False,
+        prior_probability=0.01,
+        activation="relu",
+        kernel_initializer="random_normal",
+        bias_initializer="zeros",
+        kernel_regularizer=None,
+        bias_regularizer=None,
+        use_group_norm=False,
+        data_format=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.output_filters = output_filters
+        self.num_filters = num_filters
+        self.num_conv_layers = num_conv_layers
+        self.use_prior_probability = use_prior_probability
+        self.prior_probability = prior_probability
+        self.activation = keras.activations.get(activation)
+        self.kernel_initializer = keras.initializers.get(kernel_initializer)
+        self.bias_initializer = keras.initializers.get(bias_initializer)
+        if kernel_regularizer is not None:
+            self.kernel_regularizer = keras.regularizers.get(kernel_regularizer)
+        else:
+            self.kernel_regularizer = None
+        if bias_regularizer is not None:
+            self.bias_regularizer = keras.regularizers.get(bias_regularizer)
+        else:
+            self.bias_regularizer = None
+        self.use_group_norm = use_group_norm
+        self.data_format = standardize_data_format(data_format)
+
+    def build(self, input_shape):
+        intermediate_shape = input_shape
+        self.conv_layers = []
+        self.group_norm_layers = []
+        for idx in range(self.num_conv_layers):
+            conv = keras.layers.Conv2D(
+                self.num_filters,
+                kernel_size=3,
+                padding="same",
+                kernel_initializer=self.kernel_initializer,
+                bias_initializer=self.bias_initializer,
+                use_bias=not self.use_group_norm,
+                kernel_regularizer=self.kernel_regularizer,
+                bias_regularizer=self.bias_regularizer,
+                data_format=self.data_format,
+                dtype=self.dtype_policy,
+                name=f"conv2d_{idx}",
+            )
+            conv.build(intermediate_shape)
+            self.conv_layers.append(conv)
+            intermediate_shape = (
+                input_shape[:-1] + (self.num_filters,)
+                if self.data_format == "channels_last"
+                else (input_shape[0], self.num_filters) + (input_shape[1:-1])
+            )
+            if self.use_group_norm:
+                group_norm = keras.layers.GroupNormalization(
+                    groups=32,
+                    axis=-1 if self.data_format == "channels_last" else 1,
+                    dtype=self.dtype_policy,
+                    name=f"group_norm_{idx}",
+                )
+                group_norm.build(intermediate_shape)
+                self.group_norm_layers.append(group_norm)
+        prior_probability = keras.initializers.Constant(
+            -1
+            * keras.ops.log(
+                (1 - self.prior_probability) / self.prior_probability
+            )
+        )
+        self.prediction_layer = keras.layers.Conv2D(
+            self.output_filters,
+            kernel_size=3,
+            strides=1,
+            padding="same",
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=(
+                prior_probability
+                if self.use_prior_probability
+                else self.bias_initializer
+            ),
+            kernel_regularizer=self.kernel_regularizer,
+            bias_regularizer=self.bias_regularizer,
+            dtype=self.dtype_policy,
+            name="logits_layer",
+        )
+        self.prediction_layer.build(
+            (None, None, None, self.num_filters)
+            if self.data_format == "channels_last"
+            else (None, self.num_filters, None, None)
+        )
+        self.built = True
+
+    def call(self, input):
+        x = input
+        for idx in range(self.num_conv_layers):
+            x = self.conv_layers[idx](x)
+            if self.use_group_norm:
+                x = self.group_norm_layers[idx](x)
+            x = self.activation(x)
+
+        output = self.prediction_layer(x)
+        return output
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "output_filters": self.output_filters,
+                "num_filters": self.num_filters,
+                "num_conv_layers": self.num_conv_layers,
+                "use_group_norm": self.use_group_norm,
+                "use_prior_probability": self.use_prior_probability,
+                "prior_probability": self.prior_probability,
+                "activation": keras.activations.serialize(self.activation),
+                "kernel_initializer": keras.initializers.serialize(
+                    self.kernel_initializer
+                ),
+                "bias_initializer": keras.initializers.serialize(
+                    self.kernel_initializer
+                ),
+                "kernel_regularizer": (
+                    keras.regularizers.serialize(self.kernel_regularizer)
+                    if self.kernel_regularizer is not None
+                    else None
+                ),
+                "bias_regularizer": (
+                    keras.regularizers.serialize(self.bias_regularizer)
+                    if self.bias_regularizer is not None
+                    else None
+                ),
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return (
+            input_shape[:-1] + (self.output_filters,)
+            if self.data_format == "channels_last"
+            else (input_shape[0],) + (self.output_filters,) + input_shape[1:-1]
+        )

keras_hub/src/models/retinanet/retinanet_backbone.py

@@ -0,0 +1,146 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.feature_pyramid_backbone import FeaturePyramidBackbone
+from keras_hub.src.models.retinanet.feature_pyramid import FeaturePyramid
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
+
+@keras_hub_export("keras_hub.models.RetinaNetBackbone")
+class RetinaNetBackbone(FeaturePyramidBackbone):
+    """RetinaNet Backbone.
+
+    Combines a CNN backbone (e.g., ResNet, MobileNet) with a feature pyramid
+    network (FPN)to extract multi-scale features for object detection.
+
+    Args:
+        image_encoder: `keras.Model`. The backbone model (e.g., ResNet50,
+            MobileNetV2) used to extract features from the input image.
+            It should have pyramid outputs (i.e., a dictionary mapping level
+            names like `"P2"`, `"P3"`, etc. to their corresponding feature
+            tensors).
+        min_level: int. The minimum level of the feature pyramid (e.g., 3).
+            This determines the coarsest level of features used.
+        max_level: int. The maximum level of the feature pyramid (e.g., 7).
+            This determines the finest level of features used.
+        use_p5: bool. Determines the input source for creating coarser
+            feature pyramid levels. If `True`, the output of the last backbone
+            layer (typically `'P5'` in an FPN) is used as input to create
+            higher-level feature maps (e.g., `'P6'`, `'P7'`) through
+            additional convolutional layers. If `False`, the original `'P5'`
+            feature map from the backbone is directly used as input for
+            creating the coarser levels, bypassing any further processing of
+            `'P5'` within the feature pyramid. Defaults to `False`.
+        use_fpn_batch_norm: bool. Whether to use batch normalization in the
+            feature pyramid network. Defaults to `False`.
+        image_shape: tuple. tuple. The shape of the input image (H, W, C).
+            The height and width can be `None` if they are variable.
+        data_format: str. The data format of the input image
+            (channels_first or channels_last).
+        dtype: str. The data type of the input image.
+        **kwargs: Additional keyword arguments passed to the base class.
+
+    Raises:
+        ValueError: If `min_level` is greater than `max_level`.
+        ValueError: If `backbone_max_level` is less than 5 and `max_level` is greater than or equal to 5.
+    """
+
+    def __init__(
+        self,
+        image_encoder,
+        min_level,
+        max_level,
+        use_p5,
+        use_fpn_batch_norm=False,
+        image_shape=(None, None, 3),
+        data_format=None,
+        dtype=None,
+        **kwargs,
+    ):
+
+        # === Layers ===
+        if min_level > max_level:
+            raise ValueError(
+                f"Minimum level ({min_level}) must be less than or equal to "
+                f"maximum level ({max_level})."
+            )
+
+        data_format = standardize_data_format(data_format)
+        input_levels = [
+            int(level[1]) for level in image_encoder.pyramid_outputs
+        ]
+        backbone_max_level = min(max(input_levels), max_level)
+
+        if backbone_max_level < 5 and max_level >= 5:
+            raise ValueError(
+                f"Backbone maximum level ({backbone_max_level}) is less than "
+                f"the desired maximum level ({max_level}). "
+                f"Please ensure that the backbone can generate features up to "
+                f"the specified maximum level."
+            )
+        feature_extractor = keras.Model(
+            inputs=image_encoder.inputs,
+            outputs={
+                f"P{level}": image_encoder.pyramid_outputs[f"P{level}"]
+                for level in range(min_level, backbone_max_level + 1)
+            },
+            name="backbone",
+        )
+
+        feature_pyramid = FeaturePyramid(
+            min_level=min_level,
+            max_level=max_level,
+            use_p5=use_p5,
+            name="fpn",
+            dtype=dtype,
+            data_format=data_format,
+            use_batch_norm=use_fpn_batch_norm,
+        )
+
+        # === Functional model ===
+        image_input = keras.layers.Input(image_shape, name="inputs")
+        feature_extractor_outputs = feature_extractor(image_input)
+        feature_pyramid_outputs = feature_pyramid(feature_extractor_outputs)
+
+        super().__init__(
+            inputs=image_input,
+            outputs=feature_pyramid_outputs,
+            dtype=dtype,
+            **kwargs,
+        )
+
+        # === config ===
+        self.min_level = min_level
+        self.max_level = max_level
+        self.use_p5 = use_p5
+        self.use_fpn_batch_norm = use_fpn_batch_norm
+        self.image_encoder = image_encoder
+        self.feature_pyramid = feature_pyramid
+        self.image_shape = image_shape
+        self.pyramid_outputs = feature_pyramid_outputs
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "image_encoder": keras.layers.serialize(self.image_encoder),
+                "min_level": self.min_level,
+                "max_level": self.max_level,
+                "use_p5": self.use_p5,
+                "use_fpn_batch_norm": self.use_fpn_batch_norm,
+                "image_shape": self.image_shape,
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        config.update(
+            {
+                "image_encoder": keras.layers.deserialize(
+                    config["image_encoder"]
+                ),
+            }
+        )
+
+        return super().from_config(config)

keras_hub/src/models/retinanet/retinanet_image_converter.py

@@ -0,0 +1,53 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
+from keras_hub.src.models.retinanet.retinanet_backbone import RetinaNetBackbone
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+
+@keras_hub_export("keras_hub.layers.RetinaNetImageConverter")
+class RetinaNetImageConverter(ImageConverter):
+    backbone_cls = RetinaNetBackbone
+
+    def __init__(
+        self,
+        image_size=None,
+        scale=None,
+        offset=None,
+        norm_mean=[0.485, 0.456, 0.406],
+        norm_std=[0.229, 0.224, 0.225],
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.image_size = image_size
+        self.scale = scale
+        self.offset = offset
+        self.norm_mean = norm_mean
+        self.norm_std = norm_std
+        self.built = True
+
+    @preprocessing_function
+    def call(self, inputs):
+        # TODO: https://github.com/keras-team/keras-hub/issues/1965
+        x = inputs
+        # Rescaling Image
+        if self.scale is not None:
+            x = x * self._expand_non_channel_dims(self.scale, x)
+        if self.offset is not None:
+            x = x + self._expand_non_channel_dims(self.offset, x)
+        # By default normalize using imagenet mean and std
+        if self.norm_mean:
+            x = x - self._expand_non_channel_dims(self.norm_mean, x)
+        if self.norm_std:
+            x = x / self._expand_non_channel_dims(self.norm_std, x)
+
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "norm_mean": self.norm_mean,
+                "norm_std": self.norm_std,
+            }
+        )
+        return config