[WIP] PARSeq Model

keras_hub/api/layers/__init__.py

@@ -102,6 +102,9 @@
 from keras_hub.src.models.pali_gemma.pali_gemma_image_converter import (
     PaliGemmaImageConverter as PaliGemmaImageConverter,
 )
+from keras_hub.src.models.parseq.parseq_image_converter import (
+    PARSeqImageConverter as PARSeqImageConverter,
+)
 from keras_hub.src.models.resnet.resnet_image_converter import (
     ResNetImageConverter as ResNetImageConverter,
 )

keras_hub/api/models/__init__.py

@@ -418,6 +418,18 @@
 from keras_hub.src.models.pali_gemma.pali_gemma_tokenizer import (
     PaliGemmaTokenizer as PaliGemmaTokenizer,
 )
+from keras_hub.src.models.parseq.parseq_backbone import (
+    PARSeqBackbone as PARSeqBackbone,
+)
+from keras_hub.src.models.parseq.parseq_causal_lm import (
+    PARSeqCausalLM as PARSeqCausalLM,
+)
+from keras_hub.src.models.parseq.parseq_causal_lm_preprocessor import (
+    PARSeqCausalLMPreprocessor as PARSeqCausalLMPreprocessor,
+)
+from keras_hub.src.models.parseq.parseq_tokenizer import (
+    PARSeqTokenizer as PARSeqTokenizer,
+)
 from keras_hub.src.models.phi3.phi3_backbone import Phi3Backbone as Phi3Backbone
 from keras_hub.src.models.phi3.phi3_causal_lm import (
     Phi3CausalLM as Phi3CausalLM,

keras_hub/api/tokenizers/__init__.py

@@ -65,6 +65,9 @@
 from keras_hub.src.models.pali_gemma.pali_gemma_tokenizer import (
     PaliGemmaTokenizer as PaliGemmaTokenizer,
 )
+from keras_hub.src.models.parseq.parseq_tokenizer import (
+    PARSeqTokenizer as PARSeqTokenizer,
+)
 from keras_hub.src.models.phi3.phi3_tokenizer import (
     Phi3Tokenizer as Phi3Tokenizer,
 )

keras_hub/src/models/parseq/parseq_backbone.py

@@ -0,0 +1,132 @@
+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. Defaults to `0.1`.
+        attention_dropout: float. The dropout rate for the attention weights.
+        Defaults to `0.1`.
+        dtype: str. The dtype used for layers.
+        **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)

keras_hub/src/models/parseq/parseq_backbone_test.py

@@ -0,0 +1,107 @@
+import keras
+import pytest
+from keras import ops
+
+from keras_hub.src.models.parseq.parseq_backbone import PARSeqBackbone
+from keras_hub.src.models.vit.vit_backbone import ViTBackbone
+from keras_hub.src.tests.test_case import TestCase
+
+
+class PARSeqBackboneTest(TestCase):
+    def setUp(self):
+        self.batch_size = 2
+        self.image_height = 32
+        self.image_width = 128
+        self.num_channels = 3
+
+        # Image Encoder parameters (as per your example)
+        self.vit_patch_size = (4, 8)
+        self.vit_num_layers = 2
+        self.vit_num_heads = 2
+        self.vit_hidden_dim = 64
+        self.vit_mlp_dim = self.vit_hidden_dim * 4
+
+        # PARSeq Backbone parameters
+        self.vocabulary_size = 97
+        self.max_label_length = 25
+        self.decoder_hidden_dim = self.vit_hidden_dim
+        self.num_decoder_layers = 1
+        self.num_decoder_heads = 2
+        self.decoder_mlp_dim = self.decoder_hidden_dim * 4
+
+        # Instantiate the actual ViTBackbone to be used as the image_encoder
+        self.image_encoder = ViTBackbone(
+            image_shape=(
+                self.image_height,
+                self.image_width,
+                self.num_channels,
+            ),
+            patch_size=self.vit_patch_size,
+            num_layers=self.vit_num_layers,
+            num_heads=self.vit_num_heads,
+            hidden_dim=self.vit_hidden_dim,
+            mlp_dim=self.vit_mlp_dim,
+            use_class_token=False,
+            name="image_encoder",
+        )
+
+        self.init_kwargs = {
+            "image_encoder": 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": 0.0,
+            "attention_dropout": 0.0,
+        }
+
+        # Dummy input data
+        dummy_images = keras.random.normal(
+            shape=(
+                self.batch_size,
+                self.image_height,
+                self.image_width,
+                self.num_channels,
+            ),
+        )
+
+        dummy_token_ids = keras.random.randint(
+            minval=0,
+            maxval=self.vocabulary_size,
+            shape=(self.batch_size, self.max_label_length),
+        )
+        dummy_padding_mask = ops.ones(
+            shape=(self.batch_size, self.max_label_length), dtype="int32"
+        )
+
+        self.input_data = {
+            "images": dummy_images,
+            "token_ids": dummy_token_ids,
+            "padding_mask": dummy_padding_mask,
+        }
+
+    def test_backbone_basics(self):
+        expected_shape_full = (
+            self.batch_size,
+            self.max_label_length,
+            self.vocabulary_size - 2,
+        )
+
+        self.run_backbone_test(
+            cls=PARSeqBackbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            expected_output_shape=expected_shape_full,
+            # we have image_encoder as init_kwargs which is also a backbone
+            run_quantization_check=False,
+        )
+
+    @pytest.mark.large
+    def test_saved_model(self):
+        self.run_model_saving_test(
+            cls=PARSeqBackbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )