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

keras_hub/src/models/hgnetv2/hgnetv2_backbone.py

@@ -157,7 +157,10 @@ class HGNetV2Backbone(Backbone):
             if stage_name in self.out_features
         }
         super().__init__(
-            inputs=pixel_values, outputs=feature_maps_output, **kwargs
+            inputs=pixel_values,
+            outputs=feature_maps_output,
+            dtype=dtype,
+            **kwargs,
         )
 
         # === Config ===

keras_hub/src/models/hgnetv2/hgnetv2_encoder.py

@@ -56,9 +56,10 @@ class HGNetV2Encoder(keras.layers.Layer):
         use_learnable_affine_block,
         data_format=None,
         channel_axis=None,
+        dtype=None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(dtype=dtype, **kwargs)
         self.stage_in_channels = stage_in_channels
         self.stage_mid_channels = stage_mid_channels
         self.stage_out_channels = stage_out_channels
@@ -90,7 +91,7 @@ class HGNetV2Encoder(keras.layers.Layer):
                 name=f"{self.name}_stage_{stage_idx}"
                 if self.name
                 else f"stage_{stage_idx}",
-                dtype=self.dtype,
+                dtype=dtype,
             )
             self.stages_list.append(stage_layer)
 

keras_hub/src/models/hgnetv2/hgnetv2_layers.py

@@ -17,8 +17,8 @@ class HGNetV2LearnableAffineBlock(keras.layers.Layer):
         **kwargs: Additional keyword arguments passed to the parent class.
     """
 
-    def __init__(self, scale_value=1.0, bias_value=0.0, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, scale_value=1.0, bias_value=0.0, dtype=None, **kwargs):
+        super().__init__(dtype=dtype, **kwargs)
         self.scale_value = scale_value
         self.bias_value = bias_value
 
@@ -87,9 +87,10 @@ class HGNetV2ConvLayer(keras.layers.Layer):
         use_learnable_affine_block=False,
         data_format=None,
         channel_axis=None,
+        dtype=None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(dtype=dtype, **kwargs)
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.kernel_size = kernel_size
@@ -104,6 +105,7 @@ class HGNetV2ConvLayer(keras.layers.Layer):
             padding=((pad, pad), (pad, pad)),
             data_format=self.data_format,
             name=f"{self.name}_pad" if self.name else None,
+            dtype=self.dtype_policy,
         )
         self.convolution = keras.layers.Conv2D(
             filters=self.out_channels,
@@ -156,7 +158,8 @@ class HGNetV2ConvLayer(keras.layers.Layer):
             )
         else:
             self.lab = keras.layers.Identity(
-                name=f"{self.name}_identity_lab" if self.name else None
+                name=f"{self.name}_identity_lab" if self.name else None,
+                dtype=self.dtype_policy,
             )
 
     def build(self, input_shape):
@@ -230,9 +233,10 @@ class HGNetV2ConvLayerLight(keras.layers.Layer):
         use_learnable_affine_block=False,
         data_format=None,
         channel_axis=None,
+        dtype=None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(dtype=dtype, **kwargs)
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.kernel_size = kernel_size
@@ -327,9 +331,10 @@ class HGNetV2Embeddings(keras.layers.Layer):
         use_learnable_affine_block,
         data_format=None,
         channel_axis=None,
+        dtype=None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(dtype=dtype, **kwargs)
         self.stem_channels = stem_channels
         self.hidden_act = hidden_act
         self.use_learnable_affine_block = use_learnable_affine_block
@@ -352,6 +357,7 @@ class HGNetV2Embeddings(keras.layers.Layer):
             padding=((0, 1), (0, 1)),
             data_format=self.data_format,
             name=f"{self.name}_padding1" if self.name else "padding1",
+            dtype=self.dtype_policy,
         )
         self.stem2a_layer = HGNetV2ConvLayer(
             in_channels=self.stem_channels[1],
@@ -370,6 +376,7 @@ class HGNetV2Embeddings(keras.layers.Layer):
             padding=((0, 1), (0, 1)),
             data_format=self.data_format,
             name=f"{self.name}_padding2" if self.name else "padding2",
+            dtype=self.dtype_policy,
         )
         self.stem2b_layer = HGNetV2ConvLayer(
             in_channels=self.stem_channels[1] // 2,
@@ -390,10 +397,12 @@ class HGNetV2Embeddings(keras.layers.Layer):
             padding="valid",
             data_format=self.data_format,
             name=f"{self.name}_pool" if self.name else "pool",
+            dtype=self.dtype_policy,
         )
         self.concatenate_layer = keras.layers.Concatenate(
             axis=self.channel_axis,
             name=f"{self.name}_concat" if self.name else "concat",
+            dtype=self.dtype_policy,
         )
         self.stem3_layer = HGNetV2ConvLayer(
             in_channels=self.stem_channels[1] * 2,
@@ -550,9 +559,10 @@ class HGNetV2BasicLayer(keras.layers.Layer):
         use_learnable_affine_block=False,
         data_format=None,
         channel_axis=None,
+        dtype=None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(dtype=dtype, **kwargs)
         self.in_channels_arg = in_channels
         self.middle_channels = middle_channels
         self.out_channels = out_channels
@@ -635,23 +645,27 @@ class HGNetV2BasicLayer(keras.layers.Layer):
                 self.drop_path_rate,
                 noise_shape=(None, 1, 1, 1),
                 name=f"{self.name}_drop_path" if self.name else "drop_path",
+                dtype=self.dtype_policy,
             )
         else:
             self.drop_path_layer = keras.layers.Identity(
                 name=f"{self.name}_identity_drop_path"
                 if self.name
-                else "identity_drop_path"
+                else "identity_drop_path",
+                dtype=self.dtype_policy,
             )
 
         self.concatenate_layer = keras.layers.Concatenate(
             axis=self.channel_axis,
             name=f"{self.name}_concat" if self.name else "concat",
+            dtype=self.dtype_policy,
         )
         if self.residual:
             self.add_layer = keras.layers.Add(
                 name=f"{self.name}_add_residual"
                 if self.name
-                else "add_residual"
+                else "add_residual",
+                dtype=self.dtype_policy,
             )
 
     def build(self, input_shape):
@@ -794,9 +808,10 @@ class HGNetV2Stage(keras.layers.Layer):
         drop_path: float = 0.0,
         data_format=None,
         channel_axis=None,
+        dtype=None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(dtype=dtype, **kwargs)
         self.stage_in_channels = stage_in_channels
         self.stage_mid_channels = stage_mid_channels
         self.stage_out_channels = stage_out_channels
@@ -842,7 +857,8 @@ class HGNetV2Stage(keras.layers.Layer):
             self.downsample_layer = keras.layers.Identity(
                 name=f"{self.name}_identity_downsample"
                 if self.name
-                else "identity_downsample"
+                else "identity_downsample",
+                dtype=self.dtype_policy,
             )
 
         self.blocks_list = []

keras_hub/src/models/parseq/parseq_backbone.py

@@ -0,0 +1,134 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.models.parseq.parseq_decoder import PARSeqDecoder
+
+
+@keras_hub_export("keras_hub.models.PARSeqBackbone")
+class PARSeqBackbone(Backbone):
+    """Scene Text Detection with PARSeq.
+
+    Performs OCR in natural scenes using the PARSeq model described in [Scene
+    Text Recognition with Permuted Autoregressive Sequence Models](
+    https://arxiv.org/abs/2207.06966). PARSeq is a ViT-based model that allows
+    iterative decoding by performing an autoregressive decoding phase, followed
+    by a refinement phase.
+
+    Args:
+        image_encoder: keras.Model. The image encoder model.
+        vocabulary_size: int. The size of the vocabulary.
+        max_label_length: int. The maximum length of the label sequence.
+        decoder_hidden_dim: int. The dimension of the decoder hidden layers.
+        num_decoder_layers: int. The number of decoder layers.
+        num_decoder_heads: int. The number of attention heads in the decoder.
+        decoder_mlp_dim: int. The dimension of the decoder MLP hidden layer.
+        dropout_rate: float. The dropout rate for the decoder network.
+            Defaults to `0.1`.
+        attention_dropout: float. The dropout rate for the attention weights.
+            Defaults to `0.1`.
+        dtype: str. `None`, str, or `keras.mixed_precision.DTypePolicy`. The
+            dtype to use for the computations and weights.
+        **kwargs: Additional keyword arguments passed to the base
+            `keras.Model` constructor.
+    """
+
+    def __init__(
+        self,
+        image_encoder,
+        vocabulary_size,
+        max_label_length,
+        decoder_hidden_dim,
+        num_decoder_layers,
+        num_decoder_heads,
+        decoder_mlp_dim,
+        dropout_rate=0.1,
+        attention_dropout=0.1,
+        dtype=None,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.image_encoder = image_encoder
+        self.decoder = PARSeqDecoder(
+            vocabulary_size=vocabulary_size,
+            max_label_length=max_label_length,
+            num_layers=num_decoder_layers,
+            num_heads=num_decoder_heads,
+            hidden_dim=decoder_hidden_dim,
+            mlp_dim=decoder_mlp_dim,
+            dropout_rate=dropout_rate,
+            attention_dropout=attention_dropout,
+            name="decoder",
+            dtype=dtype,
+        )
+        self.head = keras.layers.Dense(
+            vocabulary_size - 2,  # We don't predict <bos> nor <pad>
+            dtype=dtype,
+        )
+
+        # === Functional Model ===
+        image_input = self.image_encoder.input
+
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+
+        memory = self.image_encoder(image_input)
+        target_out = self.decoder(
+            token_id_input, memory, padding_mask=padding_mask_input
+        )
+        logits = self.head(target_out)
+
+        # === Config ===
+        self.vocabulary_size = vocabulary_size
+        self.max_label_length = max_label_length
+        self.decoder_hidden_dim = decoder_hidden_dim
+        self.num_decoder_layers = num_decoder_layers
+        self.num_decoder_heads = num_decoder_heads
+        self.decoder_mlp_dim = decoder_mlp_dim
+        self.dropout_rate = dropout_rate
+        self.attention_dropout = attention_dropout
+
+        super().__init__(
+            inputs={
+                "images": image_input,
+                "token_ids": token_id_input,
+                "padding_mask": padding_mask_input,
+            },
+            outputs=logits,
+            dtype=dtype,
+            **kwargs,
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "image_encoder": keras.layers.serialize(self.image_encoder),
+                "vocabulary_size": self.vocabulary_size,
+                "max_label_length": self.max_label_length,
+                "decoder_hidden_dim": self.decoder_hidden_dim,
+                "num_decoder_layers": self.num_decoder_layers,
+                "num_decoder_heads": self.num_decoder_heads,
+                "decoder_mlp_dim": self.decoder_mlp_dim,
+                "dropout_rate": self.dropout_rate,
+                "attention_dropout": self.attention_dropout,
+            }
+        )
+
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        config.update(
+            {
+                "image_encoder": keras.layers.deserialize(
+                    config["image_encoder"]
+                ),
+            }
+        )
+
+        return super().from_config(config)