keras-hub-nightly 0.23.0.dev202508260411__py3-none-any.whl → 0.23.0.dev202508280418__py3-none-any.whl

keras_hub/src/models/parseq/parseq_decoder.py

@@ -0,0 +1,418 @@
+import keras
+from keras import ops
+
+from keras_hub.src.layers.modeling.cached_multi_head_attention import (
+    CachedMultiHeadAttention,
+)
+from keras_hub.src.layers.modeling.transformer_layer_utils import (
+    compute_causal_mask,
+)
+from keras_hub.src.layers.modeling.transformer_layer_utils import (
+    merge_padding_and_attention_mask,
+)
+from keras_hub.src.models.vit.vit_layers import MLP
+
+
+class PARSeqDecoderBlock(keras.layers.Layer):
+    """A decoder block for the PARSeq model.
+
+    This block consists of self-attention, cross-attention, and a multilayer
+    perceptron (MLP). It also includes layer normalization and dropout layers.
+
+    Args:
+        hidden_dim: int. The dimension of the hidden layers.
+        num_heads: int. The number of attention heads.
+        mlp_dim: int. The dimension of the MLP hidden layer.
+        dropout_rate: float. The dropout rate used in the feedforward layers.
+        attention_dropout: float. The dropout rate for the attention weights.
+        layer_norm_epsilon: float. A small float added to the denominator for
+            numerical stability in layer normalization.
+        **kwargs: Additional keyword arguments passed to the base
+            `keras.layers.Layer` constructor.
+    """
+
+    def __init__(
+        self,
+        hidden_dim,
+        num_heads,
+        mlp_dim,
+        dropout_rate=0.1,
+        attention_dropout=0.1,
+        layer_norm_epsilon=1e-5,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        key_dim = hidden_dim // num_heads
+
+        # === Config ===
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.mlp_dim = mlp_dim
+        self.key_dim = key_dim
+        self.dropout_rate = dropout_rate
+        self.attention_dropout = attention_dropout
+        self.layer_norm_epsilon = layer_norm_epsilon
+
+    def build(self, input_shape):
+        self.query_layer_norm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            name="query_layer_norm",
+            dtype=self.dtype_policy,
+        )
+        self.query_layer_norm.build(input_shape)
+        self.content_layer_norm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            name="content_layer_norm",
+            dtype=self.dtype_policy,
+        )
+        self.content_layer_norm.build(input_shape)
+        self.self_attention = CachedMultiHeadAttention(
+            num_heads=self.num_heads,
+            key_dim=self.key_dim,
+            dropout=self.attention_dropout,
+            name="self_attention",
+            dtype=self.dtype_policy,
+        )
+        self.self_attention.build(input_shape, input_shape)
+        self.cross_attention = CachedMultiHeadAttention(
+            num_heads=self.num_heads,
+            key_dim=self.key_dim,
+            dropout=self.attention_dropout,
+            name="cross_attention",
+            dtype=self.dtype_policy,
+        )
+        self.cross_attention.build(input_shape, input_shape)
+
+        self.layer_norm_1 = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            name="ln_1",
+            dtype=self.dtype_policy,
+        )
+        self.layer_norm_1.build((None, None, self.hidden_dim))
+        self.layer_norm_2 = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            name="ln_2",
+            dtype=self.dtype_policy,
+        )
+        self.layer_norm_2.build((None, None, self.hidden_dim))
+        self.mlp = MLP(
+            hidden_dim=self.hidden_dim,
+            mlp_dim=self.mlp_dim,
+            dropout_rate=self.dropout_rate,
+            name="mlp",
+            dtype=self.dtype_policy,
+        )
+        self.mlp.build((None, None, self.hidden_dim))
+        self.dropout = keras.layers.Dropout(
+            rate=self.dropout_rate,
+            dtype=self.dtype_policy,
+            name="decoder_block_dropout",
+        )
+
+        self.built = True
+
+    def forward_stream(
+        self,
+        target,
+        target_norm,
+        target_kv,
+        memory,
+        padding_mask=None,
+        self_attention_cache=None,
+        self_attention_cache_update_index=0,
+        train_attention_mask=None,
+    ):
+        self_attention_new_cache = None
+        if train_attention_mask is None:
+            target_attention_mask = self._compute_attention_mask(
+                target_norm,
+                padding_mask,
+                self_attention_cache,
+                self_attention_cache_update_index,
+            )
+        else:
+            target_attention_mask = merge_padding_and_attention_mask(
+                target_norm, padding_mask, attention_mask=train_attention_mask
+            )
+
+        if self_attention_cache is not None:
+            target2, self_attention_new_cache = self.self_attention(
+                target_norm,
+                target_kv,
+                target_kv,
+                attention_mask=target_attention_mask,
+                cache=self_attention_cache,
+                cache_update_index=self_attention_cache_update_index,
+            )
+        else:
+            target2 = self.self_attention(
+                target_norm,
+                target_kv,
+                target_kv,
+                attention_mask=target_attention_mask,
+            )
+        target = ops.add(target, self.dropout(target2))
+        target2 = self.cross_attention(
+            self.layer_norm_1(target),
+            memory,
+            memory,
+        )
+        target = ops.add(target, self.dropout(target2))
+
+        target2 = self.mlp(self.layer_norm_2(target))
+        target = ops.add(target, target2)
+
+        return target, self_attention_new_cache
+
+    def call(
+        self,
+        query,
+        content,
+        memory,
+        padding_mask=None,
+        update_content=True,
+        query_self_attention_cache=None,
+        query_self_attention_cache_update_index=0,
+        content_self_attention_cache=None,
+        content_self_attention_cache_update_index=0,
+        query_mask=None,
+        content_mask=None,
+    ):
+        # position + token embeddings
+        query_norm = self.query_layer_norm(query)
+        # position embeddings
+        content_norm = self.content_layer_norm(content)
+        (
+            query,
+            query_self_attention_new_cache,
+        ) = self.forward_stream(
+            query,
+            query_norm,
+            content_norm,
+            memory,
+            padding_mask=padding_mask,
+            train_attention_mask=query_mask,
+            self_attention_cache=query_self_attention_cache,
+            self_attention_cache_update_index=query_self_attention_cache_update_index,
+        )
+
+        if update_content:
+            (
+                content,
+                content_self_attention_new_cache,
+            ) = self.forward_stream(
+                content,
+                content_norm,
+                content_norm,
+                memory,  # image embeddings (encoder embeddings)
+                padding_mask=padding_mask,
+                train_attention_mask=content_mask,
+                self_attention_cache=content_self_attention_cache,
+                self_attention_cache_update_index=content_self_attention_cache_update_index,
+            )
+
+        return_values = [query, content]
+
+        if query_self_attention_cache is not None:
+            return_values.append(query_self_attention_new_cache)
+        if update_content and content_self_attention_cache is not None:
+            return_values.append(content_self_attention_new_cache)
+        elif not update_content and content_self_attention_cache is not None:
+            return_values.append(content_self_attention_cache)
+
+        return tuple(return_values)
+
+    def _compute_attention_mask(
+        self, x, padding_mask, cache, cache_update_index
+    ):
+        decoder_mask = merge_padding_and_attention_mask(
+            inputs=x, padding_mask=padding_mask, attention_mask=None
+        )
+        batch_size = ops.shape(x)[0]
+        input_length = output_length = ops.shape(x)[1]
+        if cache is not None:
+            input_length = ops.shape(cache)[2]
+
+        causal_mask = compute_causal_mask(
+            batch_size=batch_size,
+            input_length=input_length,
+            output_length=output_length,
+            cache_index=cache_update_index,
+        )
+
+        return (
+            ops.minimum(decoder_mask, causal_mask)
+            if decoder_mask is not None
+            else causal_mask
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "key_dim": self.key_dim,
+                "mlp_dim": self.mlp_dim,
+                "dropout_rate": self.dropout_rate,
+                "attention_dropout": self.attention_dropout,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+            }
+        )
+        return config
+
+
+class PARSeqDecoder(keras.layers.Layer):
+    """The PARSeq decoder.
+
+    This decoder consists of multiple decoder blocks and a token embedding
+    layer. It takes token IDs and memory from the encoder as input and outputs a
+    sequence of hidden states.
+
+    Args:
+        vocabulary_size: int. The size of the vocabulary.
+        max_label_length: int. The maximum length of the label sequence.
+        num_layers: int. The number of decoder layers.
+        hidden_dim: int. The dimension of the hidden layers.
+        mlp_dim: int. The dimension of the MLP hidden layer.
+        num_heads: int. The number of attention heads.
+        dropout_rate: float. The dropout rate.
+        attention_dropout: float. The dropout rate for the attention weights.
+        layer_norm_epsilon: float. A small float added to the denominator for
+            numerical stability in layer normalization.
+        **kwargs: Additional keyword arguments passed to the base
+            `keras.layers.Layer` constructor.
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        max_label_length,
+        num_layers,
+        hidden_dim,
+        mlp_dim,
+        num_heads,
+        dropout_rate=0.1,
+        attention_dropout=0.1,
+        layer_norm_epsilon=1e-5,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        # === Config ===
+        self.vocabulary_size = vocabulary_size
+        self.max_label_length = max_label_length
+        self.hidden_dim = hidden_dim
+        self.mlp_dim = mlp_dim
+        self.num_heads = num_heads
+        self.dropout_rate = dropout_rate
+        self.attention_dropout = attention_dropout
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.num_layers = num_layers
+
+    def build(self, input_shape):
+        self.token_embedding = keras.layers.Embedding(
+            input_dim=self.vocabulary_size,
+            output_dim=self.hidden_dim,
+            dtype=self.dtype_policy,
+            name="token_embedding",
+        )
+        self.token_embedding.build((1, self.vocabulary_size))
+        self.pos_query_embeddings = self.add_weight(
+            shape=(1, self.max_label_length + 1, self.hidden_dim),
+            name="pos_query_embeddings",
+            dtype=self.dtype,
+        )
+        self.dropout = keras.layers.Dropout(
+            self.dropout_rate, dtype=self.dtype_policy, name="decoder_dropout"
+        )
+        self.decoder_layers = []
+        for i in range(self.num_layers):
+            decoder_layer = PARSeqDecoderBlock(
+                hidden_dim=self.hidden_dim,
+                num_heads=self.num_heads,
+                mlp_dim=self.mlp_dim,
+                dropout_rate=self.dropout_rate,
+                attention_dropout=self.attention_dropout,
+                layer_norm_epsilon=self.layer_norm_epsilon,
+                dtype=self.dtype_policy,
+                name=f"decoder_layer_{i}",
+            )
+            decoder_layer.build((None, None, self.hidden_dim))
+            self.decoder_layers.append(decoder_layer)
+
+        self.layer_norm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="layer_norm",
+        )
+        self.layer_norm.build((None, None, self.hidden_dim))
+        self.built = True
+
+    def call(
+        self,
+        token_ids,
+        memory,
+        padding_mask=None,
+        query_mask=None,
+        content_mask=None,
+    ):
+        bs, tokens_length = ops.shape(token_ids)
+        # <bos> stands for the null context. We only supply position information
+        # for characters after <bos>.
+        null_context = self.hidden_dim**0.5 * self.token_embedding(
+            token_ids[:, :1]
+        )
+        if tokens_length > 1:
+            content = self.pos_query_embeddings[:, : tokens_length - 1, :]
+            content = content + self.hidden_dim**0.5 * self.token_embedding(
+                token_ids[:, 1:]
+            )
+            content = ops.concatenate([null_context, content], axis=1)
+        else:
+            content = null_context
+
+        content = self.dropout(content)
+
+        query = ops.multiply(
+            ops.ones((bs, 1, 1), dtype=self.dtype),
+            self.pos_query_embeddings[:, :tokens_length, :],
+        )
+        query = self.dropout(query)
+
+        for i, decoder_layer in enumerate(self.decoder_layers):
+            last = i == self.num_layers - 1
+            query, content = decoder_layer(
+                query=query,
+                content=content,
+                memory=memory,
+                padding_mask=padding_mask,
+                update_content=not last,
+                query_mask=query_mask,
+                content_mask=content_mask,
+            )
+
+        query = self.layer_norm(query)
+
+        return query
+
+    def compute_output_shape(self, input_shape):
+        return (None, None, self.hidden_dim)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "max_label_length": self.max_label_length,
+                "num_layers": self.num_layers,
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "mlp_dim": self.mlp_dim,
+                "dropout_rate": self.dropout_rate,
+                "attention_dropout": self.attention_dropout,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+            }
+        )
+        return config

keras_hub/src/models/parseq/parseq_image_converter.py

@@ -0,0 +1,8 @@
+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.parseq.parseq_backbone import PARSeqBackbone
+
+
+@keras_hub_export("keras_hub.layers.PARSeqImageConverter")
+class PARSeqImageConverter(ImageConverter):
+    backbone_cls = PARSeqBackbone

keras_hub/src/models/parseq/parseq_tokenizer.py

@@ -0,0 +1,221 @@
+import os
+import re
+from typing import Iterable
+
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.tokenizers import tokenizer
+from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
+from keras_hub.src.utils.tensor_utils import is_int_dtype
+from keras_hub.src.utils.tensor_utils import is_string_dtype
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+try:
+    import tensorflow as tf
+    import tensorflow_text as tf_text
+except ImportError:
+    tf = None
+    tf_text = None
+
+PARSEQ_VOCAB = list(
+    "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"
+    "\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
+)
+
+VOCAB_FILENAME = "vocabulary.txt"
+
+
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.PARSeqTokenizer",
+        "keras_hub.models.PARSeqTokenizer",
+    ]
+)
+class PARSeqTokenizer(tokenizer.Tokenizer):
+    """A Tokenizer for PARSeq models, designed for OCR tasks.
+
+    This tokenizer converts strings into sequences of integer IDs or string
+    tokens, and vice-versa. It supports various preprocessing steps such as
+    whitespace removal, Unicode normalization, and limiting the maximum label
+    length. It also provides functionality to save and load the vocabulary
+    from a file.
+
+    Args:
+        vocabulary: str. A string or iterable representing the vocabulary to
+            use. If a string, it's treated as the path to a vocabulary file.
+            If an iterable, it's treated as a list of characters forming
+            the vocabulary. Defaults to `PARSEQ_VOCAB`.
+        remove_whitespace: bool. Whether to remove whitespace characters from
+            the input. Defaults to `True`.
+        normalize_unicode: bool. Whether to normalize Unicode characters in the
+            input using NFKD normalization and remove non-ASCII characters.
+            Defaults to `True`.
+        max_label_length: int. The maximum length of the tokenized output.
+            Longer labels will be truncated. Defaults to `25`.
+        dtype: str. The data type of the tokenized output. Must be an integer
+            type (e.g., "int32") or a string type ("string").
+            Defaults to `"int32"`.
+        **kwargs: Additional keyword arguments passed to the base
+            `keras.layers.Layer` constructor.
+    """
+
+    def __init__(
+        self,
+        vocabulary=PARSEQ_VOCAB,
+        remove_whitespace=True,
+        normalize_unicode=True,
+        max_label_length=25,
+        dtype="int32",
+        **kwargs,
+    ):
+        if not is_int_dtype(dtype) and not is_string_dtype(dtype):
+            raise ValueError(
+                "Output dtype must be an integer type or a string. "
+                f"Received: dtype={dtype}"
+            )
+        super().__init__(dtype=dtype, **kwargs)
+        self.remove_whitespace = remove_whitespace
+        self.normalize_unicode = normalize_unicode
+        self.max_label_length = max_label_length
+        self.file_assets = [VOCAB_FILENAME]
+
+        self.set_vocabulary(vocabulary)
+
+    def save_assets(self, dir_path):
+        path = os.path.join(dir_path, VOCAB_FILENAME)
+        with open(path, "w", encoding="utf-8") as file:
+            for token in self.vocabulary:
+                file.write(f"{token}\n")
+
+    def load_assets(self, dir_path):
+        path = os.path.join(dir_path, VOCAB_FILENAME)
+        self.set_vocabulary(path)
+
+    def set_vocabulary(self, vocabulary):
+        """Set the tokenizer vocabulary to a file or list of strings."""
+        if vocabulary is None:
+            self.vocabulary = None
+            return
+
+        if isinstance(vocabulary, str):
+            with open(vocabulary, "r", encoding="utf-8") as file:
+                self.vocabulary = [line.rstrip() for line in file]
+                self.vocabulary = "".join(self.vocabulary)
+        elif isinstance(vocabulary, Iterable):
+            self.vocabulary = "".join(vocabulary)
+        else:
+            raise ValueError(
+                "Vocabulary must be an file path or list of terms. "
+                f"Received: vocabulary={vocabulary}"
+            )
+
+        self.lowercase_only = self.vocabulary == self.vocabulary.lower()
+        self.uppercase_only = self.vocabulary == self.vocabulary.upper()
+        escaped_charset = re.escape(self.vocabulary)  # Escape for safe regex
+        self.unsupported_regex = f"[^{escaped_charset}]"
+        self._itos = ("[E]",) + tuple(self.vocabulary) + ("[B]", "[P]")
+        self._stoi = {s: i for i, s in enumerate(self._itos)}
+
+        self._add_special_token("[B]", "start_token")
+        self._add_special_token("[E]", "end_token")
+        self._add_special_token("[P]", "pad_token")
+        # Create lookup tables.
+        self.char_to_id = tf.lookup.StaticHashTable(
+            initializer=tf.lookup.KeyValueTensorInitializer(
+                keys=list(self._stoi.keys()),
+                values=list(self._stoi.values()),
+                key_dtype=tf.string,
+                value_dtype=tf.int32,
+            ),
+            default_value=self._stoi["[E]"],
+        )
+        self.id_to_char = tf.lookup.StaticHashTable(
+            initializer=tf.lookup.KeyValueTensorInitializer(
+                keys=list(self._stoi.values()),
+                values=list(self._stoi.keys()),
+                key_dtype=tf.int32,
+                value_dtype=tf.string,
+            ),
+            default_value=self.pad_token,
+        )
+
+    def get_vocabulary(self):
+        """Get the tokenizer vocabulary as a list of strings tokens."""
+        return list(self.vocabulary)
+
+    def id_to_token(self, id):
+        if id >= self.vocabulary_size() or id < 0:
+            raise ValueError(
+                f"`id` must be in range [0, {self.vocabulary_size() - 1}]. "
+                f"Received: {id}"
+            )
+        return self._itos[id]
+
+    def token_to_id(self, token):
+        return self._stoi[token]
+
+    def _preprocess(self, inputs):
+        """Performs preprocessing include only characters from ASCII."""
+        if self.remove_whitespace:
+            inputs = tf.strings.regex_replace(inputs, r"\s+", "")
+
+        if self.normalize_unicode:
+            inputs = tf_text.normalize_utf8(inputs, normalization_form="NFKD")
+            inputs = tf.strings.regex_replace(inputs, r"[^!-~]", "")
+
+        if self.lowercase_only:
+            inputs = tf.strings.lower(inputs)
+        elif self.uppercase_only:
+            inputs = tf.strings.upper(inputs)
+
+        inputs = tf.strings.regex_replace(inputs, self.unsupported_regex, "")
+        inputs = tf.strings.substr(inputs, 0, self.max_label_length)
+
+        return inputs
+
+    @preprocessing_function
+    def tokenize(self, inputs):
+        inputs = tf.convert_to_tensor(inputs)
+        unbatched = inputs.shape.rank == 0
+        if unbatched:
+            inputs = tf.expand_dims(inputs, 0)
+
+        inputs = tf.map_fn(
+            self._preprocess, inputs, fn_output_signature=tf.string
+        )
+
+        token_ids = tf.cond(
+            tf.size(inputs) > 0,
+            lambda: self.char_to_id.lookup(
+                tf.strings.unicode_split(inputs, "UTF-8")
+            ),
+            lambda: tf.RaggedTensor.from_row_splits(
+                values=tf.constant([], dtype=tf.int32),
+                row_splits=tf.constant([0], dtype=tf.int64),
+            ),
+        )
+        if unbatched:
+            token_ids = tf.squeeze(token_ids, 0)
+            tf.ensure_shape(token_ids, shape=[self.max_label_length])
+        return token_ids
+
+    @preprocessing_function
+    def detokenize(self, inputs):
+        inputs, unbatched, rectangular = convert_to_ragged_batch(inputs)
+        # tf-text sentencepiece does not handle int64.
+        inputs = tf.cast(inputs, "int32")
+        outputs = self.id_to_char.lookup(inputs)
+        if unbatched:
+            outputs = tf.squeeze(outputs, 0)
+        return outputs
+
+    def vocabulary_size(self):
+        """Get the integer size of the tokenizer vocabulary."""
+        return len(self.vocabulary) + 3
+
+    def compute_output_spec(self, input_spec):
+        return keras.KerasTensor(
+            input_spec.shape + (self.max_label_length,),
+            dtype=self.compute_dtype,
+        )

keras_hub/src/tests/test_case.py

@@ -499,6 +499,7 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         init_kwargs,
         input_data,
         expected_output_shape,
+        spatial_output_keys=None,
         expected_pyramid_output_keys=None,
         expected_pyramid_image_sizes=None,
         variable_length_data=None,
@@ -557,12 +558,47 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
                 input_data = ops.transpose(input_data, axes=(2, 0, 1))
             elif len(input_data_shape) == 4:
                 input_data = ops.transpose(input_data, axes=(0, 3, 1, 2))
-            if len(expected_output_shape) == 3:
+            if isinstance(expected_output_shape, dict):
+                # Handle dictionary of shapes.
+                transposed_shapes = {}
+                for key, shape in expected_output_shape.items():
+                    if spatial_output_keys and key not in spatial_output_keys:
+                        transposed_shapes[key] = shape
+                        continue
+                    if len(shape) == 3:
+                        transposed_shapes[key] = (shape[0], shape[2], shape[1])
+                    elif len(shape) == 4:
+                        transposed_shapes[key] = (
+                            shape[0],
+                            shape[3],
+                            shape[1],
+                            shape[2],
+                        )
+                    else:
+                        transposed_shapes[key] = shape
+                expected_output_shape = transposed_shapes
+            elif len(expected_output_shape) == 3:
                 x = expected_output_shape
                 expected_output_shape = (x[0], x[2], x[1])
             elif len(expected_output_shape) == 4:
                 x = expected_output_shape
                 expected_output_shape = (x[0], x[3], x[1], x[2])
+            original_init_kwargs = init_kwargs.copy()
+            init_kwargs = original_init_kwargs.copy()
+            # Handle nested `keras.Model` instances passed within `init_kwargs`.
+            for k, v in init_kwargs.items():
+                if isinstance(v, keras.Model) and hasattr(v, "data_format"):
+                    config = v.get_config()
+                    config["data_format"] = "channels_first"
+                    if (
+                        "image_shape" in config
+                        and config["image_shape"] is not None
+                        and len(config["image_shape"]) == 3
+                    ):
+                        config["image_shape"] = tuple(
+                            reversed(config["image_shape"])
+                        )
+                    init_kwargs[k] = v.__class__.from_config(config)
             if "image_shape" in init_kwargs:
                 init_kwargs = init_kwargs.copy()
                 init_kwargs["image_shape"] = tuple(