keras-hub-nightly 0.16.1.dev202410020340__py3-none-any.whl → 0.19.0.dev202501260345__py3-none-any.whl

keras_hub/src/models/roberta/roberta_presets.py

@@ -5,26 +5,24 @@ backbone_presets = {
         "metadata": {
             "description": (
                 "12-layer RoBERTa model where case is maintained."
-                "Trained on English Wikipedia, BooksCorpus, CommonCraw, and OpenWebText."
+                "Trained on English Wikipedia, BooksCorpus, CommonCraw, and "
+                "OpenWebText."
             ),
             "params": 124052736,
-            "official_name": "RoBERTa",
             "path": "roberta",
-            "model_card": "https://github.com/facebookresearch/fairseq/blob/main/examples/roberta/README.md",
         },
-        "kaggle_handle": "kaggle://keras/roberta/keras/roberta_base_en/2",
+        "kaggle_handle": "kaggle://keras/roberta/keras/roberta_base_en/3",
     },
     "roberta_large_en": {
         "metadata": {
             "description": (
                 "24-layer RoBERTa model where case is maintained."
-                "Trained on English Wikipedia, BooksCorpus, CommonCraw, and OpenWebText."
+                "Trained on English Wikipedia, BooksCorpus, CommonCraw, and "
+                "OpenWebText."
             ),
             "params": 354307072,
-            "official_name": "RoBERTa",
             "path": "roberta",
-            "model_card": "https://github.com/facebookresearch/fairseq/blob/main/examples/roberta/README.md",
         },
-        "kaggle_handle": "kaggle://keras/roberta/keras/roberta_large_en/2",
+        "kaggle_handle": "kaggle://keras/roberta/keras/roberta_large_en/3",
     },
 }

keras_hub/src/models/roberta/roberta_text_classifier.py

@@ -38,9 +38,9 @@ class RobertaTextClassifier(TextClassifier):
     Args:
         backbone: A `keras_hub.models.RobertaBackbone` instance.
         num_classes: int. Number of classes to predict.
-        preprocessor: A `keras_hub.models.RobertaTextClassifierPreprocessor` or `None`. If
-            `None`, this model will not apply preprocessing, and inputs should
-            be preprocessed before calling the model.
+        preprocessor: A `keras_hub.models.RobertaTextClassifierPreprocessor` or
+            `None`. If `None`, this model will not apply preprocessing, and
+            inputs should be preprocessed before calling the model.
         activation: Optional `str` or callable. The activation function to use
             on the model outputs. Set `activation="softmax"` to return output
             probabilities. Defaults to `None`.

keras_hub/src/models/sam/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.sam.sam_backbone import SAMBackbone
+from keras_hub.src.models.sam.sam_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, SAMBackbone)

keras_hub/src/models/sam/sam_backbone.py

@@ -9,8 +9,8 @@ class SAMBackbone(Backbone):
     """A backbone for the Segment Anything Model (SAM).
 
     Args:
-        image_encoder: `keras_hub.models.ViTDetBackbone`. A feature extractor for
-            the input images.
+        image_encoder: `keras_hub.models.ViTDetBackbone`. A feature extractor
+            for the input images.
         prompt_encoder: `keras_hub.layers.SAMPromptEncoder`. A Keras layer to
             compute embeddings for points, box, and mask prompt.
         mask_decoder: `keras_hub.layers.SAMMaskDecoder`. A Keras layer to
@@ -68,7 +68,6 @@ class SAMBackbone(Backbone):
         image_encoder=image_encoder,
         prompt_encoder=prompt_encoder,
         mask_decoder=mask_decoder,
-        image_shape=(image_size, image_size, 3),
     )
     backbone(input_data)
     ```

keras_hub/src/models/sam/sam_image_converter.py

@@ -1,10 +1,8 @@
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.preprocessing.resizing_image_converter import (
-    ResizingImageConverter,
-)
+from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
 from keras_hub.src.models.sam.sam_backbone import SAMBackbone
 
 
 @keras_hub_export("keras_hub.layers.SAMImageConverter")
-class SAMImageConverter(ResizingImageConverter):
+class SAMImageConverter(ImageConverter):
     backbone_cls = SAMBackbone

keras_hub/src/models/sam/sam_image_segmenter.py

@@ -31,7 +31,7 @@ class SAMImageSegmenter(ImageSegmenter):
 
 
     Args:
-      backbone: A `keras_hub.models.VGGBackbone` instance.
+      backbone: A `keras_hub.models.SAMBackbone` instance.
 
     Example:
     Load pretrained model using `from_preset`.
@@ -51,9 +51,9 @@ class SAMImageSegmenter(ImageSegmenter):
             (batch_size, 0, image_size, image_size, 1)
         ),
     }
-    # todo: update preset name
-    sam = keras_hub.models.SAMImageSegmenter.from_preset(`sam_base`)
-    sam(input_data)
+    sam = keras_hub.models.SAMImageSegmenter.from_preset('sam_base_sa1b')
+    outputs = sam.predict(input_data)
+    masks, iou_pred = outputs["masks"], outputs["iou_pred"]
     ```
 
     Load segment anything image segmenter with custom backbone
@@ -65,7 +65,7 @@ class SAMImageSegmenter(ImageSegmenter):
         (batch_size, image_size, image_size, 3),
         dtype="float32",
     )
-    image_encoder = ViTDetBackbone(
+    image_encoder = keras_hub.models.ViTDetBackbone(
         hidden_size=16,
         num_layers=16,
         intermediate_dim=16 * 4,
@@ -76,7 +76,7 @@ class SAMImageSegmenter(ImageSegmenter):
         window_size=2,
         image_shape=(image_size, image_size, 3),
     )
-    prompt_encoder = SAMPromptEncoder(
+    prompt_encoder = keras_hub.layers.SAMPromptEncoder(
         hidden_size=8,
         image_embedding_size=(8, 8),
         input_image_size=(
@@ -85,7 +85,7 @@ class SAMImageSegmenter(ImageSegmenter):
         ),
         mask_in_channels=16,
     )
-    mask_decoder = SAMMaskDecoder(
+    mask_decoder = keras_hub.layers.SAMMaskDecoder(
         num_layers=2,
         hidden_size=8,
         intermediate_dim=32,
@@ -95,13 +95,12 @@ class SAMImageSegmenter(ImageSegmenter):
         iou_head_depth=3,
         iou_head_hidden_dim=8,
     )
-    backbone = SAMBackbone(
+    backbone = keras_hub.models.SAMBackbone(
         image_encoder=image_encoder,
         prompt_encoder=prompt_encoder,
         mask_decoder=mask_decoder,
-        image_shape=(image_size, image_size, 3),
     )
-    sam = SAMImageSegmenter(
+    sam = keras_hub.models.SAMImageSegmenter(
         backbone=backbone
     )
     ```
@@ -115,7 +114,7 @@ class SAMImageSegmenter(ImageSegmenter):
     labels = np.array([[1., 0.]])
     box = np.array([[[[384., 384.], [640., 640.]]]])
     input_mask = np.ones((1, 1, 256, 256, 1))
-    Prepare an input dictionary:
+    # Prepare an input dictionary:
     inputs = {
         "images": image,
         "points": points,
@@ -201,17 +200,18 @@ class SAMImageSegmenter(ImageSegmenter):
     def _add_placeholder_prompts(self, inputs):
         """Adds placeholder prompt inputs for a call to SAM.
 
-        Because SAM is a functional subclass model, all inputs must be specified in
-        calls to the model. However, prompt inputs are all optional, so we have to
-        add placeholders when they're not specified by the user.
+        Because SAM is a functional subclass model, all inputs must be specified
+        in calls to the model. However, prompt inputs are all optional, so we
+        have to add placeholders when they're not specified by the user.
         """
         inputs = inputs.copy()
 
         # Get the batch shape based on the image input
         batch_size = ops.shape(inputs["images"])[0]
 
-        # The type of the placeholders must match the existing inputs with respect
-        # to whether or not they are tensors (as opposed to Numpy arrays).
+        # The type of the placeholders must match the existing inputs with
+        # respect to whether or not they are tensors (as opposed to Numpy
+        # arrays).
         zeros = ops.zeros if ops.is_tensor(inputs["images"]) else np.zeros
 
         # Fill in missing inputs.

keras_hub/src/models/sam/sam_image_segmenter_preprocessor.py

@@ -1,12 +1,22 @@
+import keras
+
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.image_segmenter_preprocessor import (
     ImageSegmenterPreprocessor,
 )
 from keras_hub.src.models.sam.sam_backbone import SAMBackbone
 from keras_hub.src.models.sam.sam_image_converter import SAMImageConverter
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 
-@keras_hub_export("keras_hub.models.SamImageSegmenterPreprocessor")
+@keras_hub_export("keras_hub.models.SAMImageSegmenterPreprocessor")
 class SAMImageSegmenterPreprocessor(ImageSegmenterPreprocessor):
     backbone_cls = SAMBackbone
     image_converter_cls = SAMImageConverter
+
+    @preprocessing_function
+    def call(self, x, y=None, sample_weight=None):
+        images = x["images"]
+        if self.image_converter:
+            x["images"] = self.image_converter(images)
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)

keras_hub/src/models/sam/sam_layers.py

@@ -170,8 +170,8 @@ class TwoWayMultiHeadAttention(keras.layers.Layer):
         key_dim: int. Size of each attention head for query, key, and
             value.
         intermediate_dim: int. Number of hidden dims to use in the mlp block.
-        skip_first_layer_pos_embedding: bool. A boolean indicating whether to skip the
-            first layer positional embeddings.
+        skip_first_layer_pos_embedding: bool. A boolean indicating whether to
+            skip the first layer positional embeddings.
         attention_downsample_rate: int, optional. The downsample rate to use
             in the attention layers. Defaults to 2.
         activation: str, optional. The activation for the mlp block's output
@@ -296,7 +296,9 @@ class TwoWayMultiHeadAttention(keras.layers.Layer):
                 "num_heads": self.num_heads,
                 "key_dim": self.key_dim,
                 "intermediate_dim": self.intermediate_dim,
-                "skip_first_layer_pos_embedding": self.skip_first_layer_pos_embedding,
+                "skip_first_layer_pos_embedding": (
+                    self.skip_first_layer_pos_embedding
+                ),
                 "attention_downsample_rate": self.attention_downsample_rate,
                 "activation": self.activation,
             }

keras_hub/src/models/sam/sam_presets.py

@@ -5,30 +5,24 @@ backbone_presets = {
         "metadata": {
             "description": ("The base SAM model trained on the SA1B dataset."),
             "params": 93735728,
-            "official_name": "SAMImageSegmenter",
             "path": "sam",
-            "model_card": "https://arxiv.org/abs/2304.02643",
         },
-        "kaggle_handle": "kaggle://kerashub/sam/keras/sam_base_sa1b/1",
+        "kaggle_handle": "kaggle://keras/sam/keras/sam_base_sa1b/5",
     },
     "sam_large_sa1b": {
         "metadata": {
             "description": ("The large SAM model trained on the SA1B dataset."),
             "params": 641090864,
-            "official_name": "SAMImageSegmenter",
             "path": "sam",
-            "model_card": "https://arxiv.org/abs/2304.02643",
         },
-        "kaggle_handle": "kaggle://kerashub/sam/keras/sam_large_sa1b/1",
+        "kaggle_handle": "kaggle://keras/sam/keras/sam_large_sa1b/5",
     },
     "sam_huge_sa1b": {
         "metadata": {
             "description": ("The huge SAM model trained on the SA1B dataset."),
             "params": 312343088,
-            "official_name": "SAMImageSegmenter",
             "path": "sam",
-            "model_card": "https://arxiv.org/abs/2304.02643",
         },
-        "kaggle_handle": "kaggle://kerashub/sam/keras/sam_huge_sa1b/1",
+        "kaggle_handle": "kaggle://keras/sam/keras/sam_huge_sa1b/5",
     },
 }

keras_hub/src/models/sam/sam_prompt_encoder.py

@@ -57,7 +57,7 @@ class SAMPromptEncoder(keras.layers.Layer):
         input_image_size=(1024, 1024),
         mask_in_channels=16,
         activation="gelu",
-        **kwargs
+        **kwargs,
     ):
         super().__init__(**kwargs)
         self.hidden_size = hidden_size
@@ -305,7 +305,9 @@ class SAMPromptEncoder(keras.layers.Layer):
         return {
             "prompt_sparse_embeddings": sparse_embeddings,
             "prompt_dense_embeddings": dense_embeddings,
-            "prompt_dense_positional_embeddings": prompt_dense_positional_embeddings,
+            "prompt_dense_positional_embeddings": (
+                prompt_dense_positional_embeddings
+            ),
         }
 
     def get_config(self):

keras_hub/src/models/sam/sam_transformer.py

@@ -31,14 +31,15 @@ class TwoWayTransformer(keras.layers.Layer):
     location and type.
 
     Args:
-        num_layers: int, optional. The num_layers of the attention blocks (the number
-            of attention blocks to use). Defaults to `2`.
+        num_layers: int, optional. The num_layers of the attention blocks
+            (the number of attention blocks to use). Defaults to `2`.
         hidden_size: int, optional. The number of features of the input image
             and point embeddings. Defaults to `256`.
         num_heads: int, optional. Number of heads to use in the attention
             layers. Defaults to `8`.
-        intermediate_dim: int, optional. The number of units in the hidden layer of
-            the MLP block used in the attention layers. Defaults to `2048`.
+        intermediate_dim: int, optional. The number of units in the hidden
+            layer of the MLP block used in the attention layers.
+            Defaults to `2048`.
         activation: str, optional. The activation of the MLP block's output
             layer used in the attention layers. Defaults to `"relu"`.
         attention_downsample_rate: int, optional. The downsample rate of the

keras_hub/src/models/segformer/__init__.py

@@ -0,0 +1,8 @@
+from keras_hub.src.models.segformer.segformer_backbone import SegFormerBackbone
+from keras_hub.src.models.segformer.segformer_image_segmenter import (
+    SegFormerImageSegmenter,
+)
+from keras_hub.src.models.segformer.segformer_presets import presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(presets, SegFormerImageSegmenter)

keras_hub/src/models/segformer/segformer_backbone.py

@@ -0,0 +1,167 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.backbone import Backbone
+
+
+@keras_hub_export("keras_hub.models.SegFormerBackbone")
+class SegFormerBackbone(Backbone):
+    """A Keras model implementing SegFormer for semantic segmentation.
+
+    This class implements the majority of the SegFormer architecture described
+    in [SegFormer: Simple and Efficient Design for Semantic Segmentation](https://arxiv.org/abs/2105.15203)
+    and based on the TensorFlow implementation
+    [from DeepVision](https://github.com/DavidLandup0/deepvision/tree/main/deepvision/models/segmentation/segformer).
+
+    SegFormers are meant to be used with the MixTransformer (MiT) encoder
+    family, and use a very lightweight all-MLP decoder head.
+
+    The MiT encoder uses a hierarchical transformer which outputs features at
+    multiple scales, similar to that of the hierarchical outputs typically
+    associated with CNNs.
+
+    Args:
+        image_encoder: `keras.Model`. The backbone network for the model that is
+            used as a feature extractor for the SegFormer encoder.
+            Should be used with the MiT backbone model
+            (`keras_hub.models.MiTBackbone`) which was created
+            specifically for SegFormers.
+        num_classes: int, the number of classes for the detection model,
+            including the background class.
+        projection_filters: int, number of filters in the
+            convolution layer projecting the concatenated features into
+            a segmentation map. Defaults to 256`.
+
+    Example:
+
+    Using the class with a custom `backbone`:
+
+    ```python
+    import keras_hub
+
+    backbone = keras_hub.models.MiTBackbone(
+        depths=[2, 2, 2, 2],
+        image_shape=(224, 224, 3),
+        hidden_dims=[32, 64, 160, 256],
+        num_layers=4,
+        blockwise_num_heads=[1, 2, 5, 8],
+        blockwise_sr_ratios=[8, 4, 2, 1],
+        max_drop_path_rate=0.1,
+        patch_sizes=[7, 3, 3, 3],
+        strides=[4, 2, 2, 2],
+    )
+
+    segformer_backbone = keras_hub.models.SegFormerBackbone(
+        image_encoder=backbone, projection_filters=256)
+    ```
+
+    Using the class with a preset `backbone`:
+
+    ```python
+    import keras_hub
+
+    backbone = keras_hub.models.MiTBackbone.from_preset("mit_b0_ade20k_512")
+    segformer_backbone = keras_hub.models.SegFormerBackbone(
+        image_encoder=backbone, projection_filters=256)
+    ```
+
+    """
+
+    def __init__(
+        self,
+        image_encoder,
+        projection_filters,
+        **kwargs,
+    ):
+        if not isinstance(image_encoder, keras.layers.Layer) or not isinstance(
+            image_encoder, keras.Model
+        ):
+            raise ValueError(
+                "Argument `image_encoder` must be a `keras.layers.Layer` "
+                f"instance or `keras.Model`. Received instead "
+                f"image_encoder={image_encoder} "
+                f"(of type {type(image_encoder)})."
+            )
+
+        # === Layers ===
+        inputs = keras.layers.Input(shape=image_encoder.input.shape[1:])
+
+        self.feature_extractor = keras.Model(
+            image_encoder.inputs, image_encoder.pyramid_outputs
+        )
+
+        features = self.feature_extractor(inputs)
+        # Get height and width of level one output
+        _, height, width, _ = features["P1"].shape
+
+        self.mlp_blocks = []
+
+        for feature_dim, feature in zip(image_encoder.hidden_dims, features):
+            self.mlp_blocks.append(
+                keras.layers.Dense(
+                    projection_filters, name=f"linear_{feature_dim}"
+                )
+            )
+
+        self.resizing = keras.layers.Resizing(
+            height, width, interpolation="bilinear"
+        )
+        self.concat = keras.layers.Concatenate(axis=-1)
+        self.linear_fuse = keras.Sequential(
+            [
+                keras.layers.Conv2D(
+                    filters=projection_filters, kernel_size=1, use_bias=False
+                ),
+                keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.9),
+                keras.layers.Activation("relu"),
+            ]
+        )
+
+        # === Functional Model ===
+        # Project all multi-level outputs onto
+        # the same dimensionality and feature map shape
+        multi_layer_outs = []
+        for index, (feature_dim, feature) in enumerate(
+            zip(image_encoder.hidden_dims, features)
+        ):
+            out = self.mlp_blocks[index](features[feature])
+            out = self.resizing(out)
+            multi_layer_outs.append(out)
+
+        # Concat now-equal feature maps
+        concatenated_outs = self.concat(multi_layer_outs[::-1])
+
+        # Fuse concatenated features into a segmentation map
+        seg = self.linear_fuse(concatenated_outs)
+
+        super().__init__(
+            inputs=inputs,
+            outputs=seg,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.projection_filters = projection_filters
+        self.image_encoder = image_encoder
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "projection_filters": self.projection_filters,
+                "image_encoder": keras.saving.serialize_keras_object(
+                    self.image_encoder
+                ),
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        if "image_encoder" in config and isinstance(
+            config["image_encoder"], dict
+        ):
+            config["image_encoder"] = keras.layers.deserialize(
+                config["image_encoder"]
+            )
+        return super().from_config(config)

keras_hub/src/models/segformer/segformer_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.segformer.segformer_backbone import SegFormerBackbone
+
+
+@keras_hub_export("keras_hub.layers.SegFormerImageConverter")
+class SegFormerImageConverter(ImageConverter):
+    backbone_cls = SegFormerBackbone

keras_hub/src/models/segformer/segformer_image_segmenter.py

@@ -0,0 +1,184 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.image_segmenter import ImageSegmenter
+from keras_hub.src.models.segformer.segformer_backbone import SegFormerBackbone
+from keras_hub.src.models.segformer.segformer_image_segmenter_preprocessor import (  # noqa: E501
+    SegFormerImageSegmenterPreprocessor,
+)
+
+
+@keras_hub_export("keras_hub.models.SegFormerImageSegmenter")
+class SegFormerImageSegmenter(ImageSegmenter):
+    """A Keras model implementing SegFormer for semantic segmentation.
+
+    This class implements the segmentation head of the SegFormer architecture
+    described in [SegFormer: Simple and Efficient Design for Semantic
+    Segmentation with Transformers] (https://arxiv.org/abs/2105.15203) and
+    [based on the TensorFlow implementation from DeepVision]
+    (https://github.com/DavidLandup0/deepvision/tree/main/deepvision/models/segmentation/segformer).
+
+    SegFormers are meant to be used with the MixTransformer (MiT) encoder
+    family, and and use a very lightweight all-MLP decoder head.
+
+    The MiT encoder uses a hierarchical transformer which outputs features at
+    multiple scales, similar to that of the hierarchical outputs typically
+    associated with CNNs.
+
+    Args:
+        image_encoder: `keras.Model`. The backbone network for the model that is
+            used as a feature extractor for the SegFormer encoder. It is
+            *intended* to be used only with the MiT backbone model
+            (`keras_hub.models.MiTBackbone`) which was created specifically for
+            SegFormers. Alternatively, can be a `keras_hub.models.Backbone` a
+            model subclassing `keras_hub.models.FeaturePyramidBackbone`, or a
+            `keras.Model` that has a `pyramid_outputs` property which is a
+            dictionary with keys "P2", "P3", "P4", and "P5" and layer names as
+            values.
+        num_classes: int, the number of classes for the detection model,
+            including the background class.
+        projection_filters: int, number of filters in the
+            convolution layer projecting the concatenated features into a
+            segmentation map. Defaults to 256`.
+
+
+    Example:
+
+    Using presets:
+
+    ```python
+    segmenter = keras_hub.models.SegFormerImageSegmenter.from_preset(
+        "segformer_b0_ade20k_512"
+    )
+
+    images = np.random.rand(1, 512, 512, 3)
+    segformer(images)
+    ```
+
+    Using the SegFormer backbone:
+
+    ```python
+    encoder = keras_hub.models.MiTBackbone.from_preset(
+        "mit_b0_ade20k_512"
+    )
+    backbone = keras_hub.models.SegFormerBackbone(
+        image_encoder=encoder,
+        projection_filters=256,
+    )
+    ```
+
+    Using the SegFormer backbone with a custom encoder:
+
+    ```python
+    images = np.ones(shape=(1, 96, 96, 3))
+    labels = np.zeros(shape=(1, 96, 96, 1))
+
+    encoder = keras_hub.models.MiTBackbone(
+        depths=[2, 2, 2, 2],
+        image_shape=(96, 96, 3),
+        hidden_dims=[32, 64, 160, 256],
+        num_layers=4,
+        blockwise_num_heads=[1, 2, 5, 8],
+        blockwise_sr_ratios=[8, 4, 2, 1],
+        max_drop_path_rate=0.1,
+        patch_sizes=[7, 3, 3, 3],
+        strides=[4, 2, 2, 2],
+    )
+
+    backbone = keras_hub.models.SegFormerBackbone(
+        image_encoder=encoder,
+        projection_filters=256,
+    )
+    segformer = keras_hub.models.SegFormerImageSegmenter(
+        backbone=backbone,
+        num_classes=4,
+    )
+    segformer(images
+    ```
+
+    Using the segmentor class with a preset backbone:
+
+    ```python
+    image_encoder = keras_hub.models.MiTBackbone.from_preset(
+        "mit_b0_ade20k_512"
+    )
+    backbone = keras_hub.models.SegFormerBackbone(
+        image_encoder=encoder,
+        projection_filters=256,
+    )
+    segformer = keras_hub.models.SegFormerImageSegmenter(
+        backbone=backbone,
+        num_classes=4,
+    )
+    ```
+    """
+
+    backbone_cls = SegFormerBackbone
+    preprocessor_cls = SegFormerImageSegmenterPreprocessor
+
+    def __init__(
+        self,
+        backbone,
+        num_classes,
+        preprocessor=None,
+        **kwargs,
+    ):
+        if not isinstance(backbone, keras.layers.Layer) or not isinstance(
+            backbone, keras.Model
+        ):
+            raise ValueError(
+                "Argument `backbone` must be a `keras.layers.Layer` instance "
+                f" or `keras.Model`. Received instead "
+                f"backbone={backbone} (of type {type(backbone)})."
+            )
+
+        # === Layers ===
+        inputs = backbone.input
+
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.dropout = keras.layers.Dropout(0.1)
+        self.output_segmentation_head = keras.layers.Conv2D(
+            filters=num_classes, kernel_size=1, strides=1
+        )
+        self.resizing = keras.layers.Resizing(
+            height=inputs.shape[1],
+            width=inputs.shape[2],
+            interpolation="bilinear",
+        )
+
+        # === Functional Model ===
+        x = self.backbone(inputs)
+        x = self.dropout(x)
+        x = self.output_segmentation_head(x)
+        output = self.resizing(x)
+
+        super().__init__(
+            inputs=inputs,
+            outputs=output,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.num_classes = num_classes
+        self.backbone = backbone
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_classes": self.num_classes,
+                "backbone": keras.saving.serialize_keras_object(self.backbone),
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        if "image_encoder" in config and isinstance(
+            config["image_encoder"], dict
+        ):
+            config["image_encoder"] = keras.layers.deserialize(
+                config["image_encoder"]
+            )
+        return super().from_config(config)