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

keras_hub/src/layers/preprocessing/image_converter.py

@@ -1,47 +1,190 @@
+import math
+
+import keras
+import numpy as np
+from keras import ops
+
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.preprocessing.preprocessing_layer import (
     PreprocessingLayer,
 )
-from keras_hub.src.utils.preset_utils import IMAGE_CONVERTER_CONFIG_FILE
+from keras_hub.src.utils.keras_utils import standardize_data_format
 from keras_hub.src.utils.preset_utils import builtin_presets
 from keras_hub.src.utils.preset_utils import find_subclass
 from keras_hub.src.utils.preset_utils import get_preset_loader
-from keras_hub.src.utils.preset_utils import save_serialized_object
+from keras_hub.src.utils.preset_utils import get_preset_saver
 from keras_hub.src.utils.python_utils import classproperty
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 
 @keras_hub_export("keras_hub.layers.ImageConverter")
 class ImageConverter(PreprocessingLayer):
-    """Convert raw image for models that support image input.
+    """Preprocess raw images into model ready inputs.
+
+    This class converts from raw images to model ready inputs. This conversion
+    proceeds in the following steps:
 
-    This class converts from raw images of any size, to preprocessed
-    images for pretrained model inputs. It is meant to be a convenient way to
-    write custom preprocessing code that is not model specific. This layer
-    should be instantiated via the `from_preset()` constructor, which will
-    create the correct subclass of this layer for the model preset.
+    1. Resize the image using to `image_size`. If `image_size` is `None`, this
+       step will be skipped.
+    2. Rescale the image by multiplying by `scale`, which can be either global
+       or per channel. If `scale` is `None`, this step will be skipped.
+    3. Offset the image by adding `offset`, which can be either global
+       or per channel. If `offset` is `None`, this step will be skipped.
 
     The layer will take as input a raw image tensor in the channels last or
     channels first format, and output a preprocessed image input for modeling.
-    The exact structure of the output will vary per model, though in most cases
-    this layer will simply resize the image to the size needed by the model
-    input.
+    This tensor can be batched (rank 4), or unbatched (rank 3).
+
+    This layer can be used with the `from_preset()` constructor to load a layer
+    that will rescale and resize an image for a specific pretrained model.
+    Using the layer this way allows writing preprocessing code that does not
+    need updating when switching between model checkpoints.
+
+    Args:
+        image_size: `(int, int)` tuple or `None`. The output size of the image,
+            not including the channels axis. If `None`, the input will not be
+            resized.
+        scale: float, tuple of floats, or `None`. The scale to apply to the
+            inputs. If `scale` is a single float, the entire input will be
+            multiplied by `scale`. If `scale` is a tuple, it's assumed to
+            contain per-channel scale value multiplied against each channel of
+            the input images. If `scale` is `None`, no scaling is applied.
+        offset: float, tuple of floats, or `None`. The offset to apply to the
+            inputs. If `offset` is a single float, the entire input will be
+            summed with `offset`. If `offset` is a tuple, it's assumed to
+            contain per-channel offset value summed against each channel of the
+            input images. If `offset` is `None`, no scaling is applied.
+        crop_to_aspect_ratio: If `True`, resize the images without aspect
+            ratio distortion. When the original aspect ratio differs
+            from the target aspect ratio, the output image will be
+            cropped so as to return the
+            largest possible window in the image (of size `(height, width)`)
+            that matches the target aspect ratio. By default
+            (`crop_to_aspect_ratio=False`), aspect ratio may not be preserved.
+        interpolation: String, the interpolation method.
+            Supports `"bilinear"`, `"nearest"`, `"bicubic"`,
+            `"lanczos3"`, `"lanczos5"`. Defaults to `"bilinear"`.
+        data_format: String, either `"channels_last"` or `"channels_first"`.
+            The ordering of the dimensions in the inputs. `"channels_last"`
+            corresponds to inputs with shape `(batch, height, width, channels)`
+            while `"channels_first"` corresponds to inputs with shape
+            `(batch, channels, height, width)`. It defaults to the
+            `image_data_format` value found in your Keras config file at
+            `~/.keras/keras.json`. If you never set it, then it will be
+            `"channels_last"`.
 
     Examples:
     ```python
-    # Resize images for `"pali_gemma_3b_224"`.
-    converter = keras_hub.layers.ImageConverter.from_preset("pali_gemma_3b_224")
-    converter(np.ones(2, 512, 512, 3)) # Output shape: (2, 224, 224, 3)
-    # Resize images for `"pali_gemma_3b_448"`.
-    converter = keras_hub.layers.ImageConverter.from_preset("pali_gemma_3b_448")
-    converter(np.ones(2, 512, 512, 3)) # Output shape: (2, 448, 448, 3)
+    # Resize raw images and scale them to [0, 1].
+    converter = keras_hub.layers.ImageConverter(
+        image_size=(128, 128),
+        scale=1. / 255,
+    )
+    converter(np.random.randint(0, 256, size=(2, 512, 512, 3)))
+
+    # Resize images to the specific size needed for a PaliGemma preset.
+    converter = keras_hub.layers.ImageConverter.from_preset(
+        "pali_gemma_3b_224"
+    )
+    converter(np.random.randint(0, 256, size=(2, 512, 512, 3)))
     ```
     """
 
     backbone_cls = None
 
+    def __init__(
+        self,
+        image_size=None,
+        scale=None,
+        offset=None,
+        crop_to_aspect_ratio=True,
+        interpolation="bilinear",
+        data_format=None,
+        **kwargs,
+    ):
+        # TODO: old arg names. Delete this block after resaving Kaggle assets.
+        if "height" in kwargs and "width" in kwargs:
+            image_size = (kwargs.pop("height"), kwargs.pop("width"))
+        if "variance" in kwargs and "mean" in kwargs:
+            std = [math.sqrt(v) for v in kwargs.pop("variance")]
+            scale = [scale / s for s in std]
+            offset = [-m / s for m, s in zip(kwargs.pop("mean"), std)]
+
+        super().__init__(**kwargs)
+
+        # Create the `Resizing` layer here even if it's not being used. That
+        # allows us to make `image_size` a settable property.
+        self.resizing = keras.layers.Resizing(
+            height=image_size[0] if image_size else None,
+            width=image_size[1] if image_size else None,
+            crop_to_aspect_ratio=crop_to_aspect_ratio,
+            interpolation=interpolation,
+            data_format=data_format,
+            dtype=self.dtype_policy,
+            name="resizing",
+        )
+        self.scale = scale
+        self.offset = offset
+        self.crop_to_aspect_ratio = crop_to_aspect_ratio
+        self.interpolation = interpolation
+        self.data_format = standardize_data_format(data_format)
+
+    @property
     def image_size(self):
-        """Returns the default size of a single image."""
-        return (None, None)
+        """Settable tuple of `(height, width)` ints. The output image shape."""
+        if self.resizing.height is None:
+            return None
+        return (self.resizing.height, self.resizing.width)
+
+    @image_size.setter
+    def image_size(self, value):
+        if value is None:
+            value = (None, None)
+        self.resizing.height = value[0]
+        self.resizing.width = value[1]
+
+    @preprocessing_function
+    def call(self, inputs):
+        x = inputs
+        if self.image_size is not None:
+            x = self.resizing(x)
+        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)
+        return x
+
+    def _expand_non_channel_dims(self, value, inputs):
+        unbatched = len(ops.shape(inputs)) == 3
+        channels_first = self.data_format == "channels_first"
+        if unbatched:
+            broadcast_dims = (1, 2) if channels_first else (0, 1)
+        else:
+            broadcast_dims = (0, 2, 3) if channels_first else (0, 1, 2)
+        # If inputs are not a tensor type, return a numpy array.
+        # This might happen when running under tf.data.
+        if ops.is_tensor(inputs):
+            # preprocessing decorator moves tensors to cpu in torch backend and
+            # processed on CPU, and then converted back to the appropriate
+            # device (potentially GPU) after preprocessing.
+            if keras.backend.backend() == "torch" and self.image_size is None:
+                return ops.expand_dims(value, broadcast_dims).cpu()
+            return ops.expand_dims(value, broadcast_dims)
+        else:
+            return np.expand_dims(value, broadcast_dims)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "image_size": self.image_size,
+                "scale": self.scale,
+                "offset": self.offset,
+                "interpolation": self.interpolation,
+                "crop_to_aspect_ratio": self.crop_to_aspect_ratio,
+            }
+        )
+        return config
 
     @classproperty
     def presets(cls):
@@ -69,13 +212,6 @@ class ImageConverter(PreprocessingLayer):
         You can run `cls.presets.keys()` to list all built-in presets available
         on the class.
 
-        This constructor can be called in one of two ways. Either from the base
-        class like `keras_hub.models.ImageConverter.from_preset()`, or from a
-        model class like
-        `keras_hub.models.PaliGemmaImageConverter.from_preset()`. If calling
-        from the base class, the subclass of the returning object will be
-        inferred from the config in the preset directory.
-
         Args:
             preset: string. A built-in preset identifier, a Kaggle Models
                 handle, a Hugging Face handle, or a path to a local directory.
@@ -85,17 +221,20 @@ class ImageConverter(PreprocessingLayer):
 
         Examples:
         ```python
+        batch = np.random.randint(0, 256, size=(2, 512, 512, 3))
+
         # Resize images for `"pali_gemma_3b_224"`.
         converter = keras_hub.layers.ImageConverter.from_preset(
             "pali_gemma_3b_224"
         )
-        converter(np.ones(2, 512, 512, 3)) # Output shape: (2, 224, 224, 3)
-        # Override arguments on the base class.
+        converter(batch) # Output shape: (2, 224, 224, 3)
+
+        # Resize images for `"pali_gemma_3b_448"` without cropping.
         converter = keras_hub.layers.ImageConverter.from_preset(
             "pali_gemma_3b_448",
             crop_to_aspect_ratio=False,
         )
-        converter(np.ones(2, 512, 512, 3)) # (2, 448, 448, 3)
+        converter(batch) # Output shape: (2, 448, 448, 3)
         ```
         """
         loader = get_preset_loader(preset)
@@ -110,8 +249,5 @@ class ImageConverter(PreprocessingLayer):
         Args:
             preset_dir: The path to the local model preset directory.
         """
-        save_serialized_object(
-            self,
-            preset_dir,
-            config_file=IMAGE_CONVERTER_CONFIG_FILE,
-        )
+        saver = get_preset_saver(preset_dir)
+        saver.save_image_converter(self)

keras_hub/src/metrics/bleu.py

@@ -164,7 +164,7 @@ class Bleu(keras.metrics.Metric):
         return inputs
 
     def _get_ngrams(self, segment, max_order):
-        """Extracts all n-grams up to a given maximum order from an input segment.
+        """Extracts all n-grams up to a given maximum order from an input.
 
         Uses Python ops. Inspired from
         https://github.com/tensorflow/nmt/blob/master/nmt/scripts/bleu.py.
@@ -329,8 +329,9 @@ class Bleu(keras.metrics.Metric):
                 return tf.squeeze(inputs, axis=-1)
             else:
                 raise ValueError(
-                    f"{tensor_name} must be of rank {base_rank}, {base_rank+1} "
-                    f"or {base_rank+2}. Found rank: {inputs.shape.rank}"
+                    f"{tensor_name} must be of rank {base_rank}, "
+                    f"{base_rank + 1}, or {base_rank + 2}. "
+                    f"Found rank: {inputs.shape.rank}"
                 )
 
         y_true = validate_and_fix_rank(y_true, "y_true", 1)

keras_hub/src/models/albert/albert_presets.py

@@ -8,11 +8,9 @@ backbone_presets = {
                 "Trained on English Wikipedia + BooksCorpus."
             ),
             "params": 11683584,
-            "official_name": "ALBERT",
             "path": "albert",
-            "model_card": "https://github.com/google-research/albert/blob/master/README.md",
         },
-        "kaggle_handle": "kaggle://keras/albert/keras/albert_base_en_uncased/2",
+        "kaggle_handle": "kaggle://keras/albert/keras/albert_base_en_uncased/5",
     },
     "albert_large_en_uncased": {
         "metadata": {
@@ -21,11 +19,9 @@ backbone_presets = {
                 "Trained on English Wikipedia + BooksCorpus."
             ),
             "params": 17683968,
-            "official_name": "ALBERT",
             "path": "albert",
-            "model_card": "https://github.com/google-research/albert/blob/master/README.md",
         },
-        "kaggle_handle": "kaggle://keras/albert/keras/albert_large_en_uncased/2",
+        "kaggle_handle": "kaggle://keras/albert/keras/albert_large_en_uncased/3",
     },
     "albert_extra_large_en_uncased": {
         "metadata": {
@@ -34,11 +30,9 @@ backbone_presets = {
                 "Trained on English Wikipedia + BooksCorpus."
             ),
             "params": 58724864,
-            "official_name": "ALBERT",
             "path": "albert",
-            "model_card": "https://github.com/google-research/albert/blob/master/README.md",
         },
-        "kaggle_handle": "kaggle://keras/albert/keras/albert_extra_large_en_uncased/2",
+        "kaggle_handle": "kaggle://keras/albert/keras/albert_extra_large_en_uncased/3",
     },
     "albert_extra_extra_large_en_uncased": {
         "metadata": {
@@ -47,10 +41,8 @@ backbone_presets = {
                 "Trained on English Wikipedia + BooksCorpus."
             ),
             "params": 222595584,
-            "official_name": "ALBERT",
             "path": "albert",
-            "model_card": "https://github.com/google-research/albert/blob/master/README.md",
         },
-        "kaggle_handle": "kaggle://keras/albert/keras/albert_extra_extra_large_en_uncased/2",
+        "kaggle_handle": "kaggle://keras/albert/keras/albert_extra_extra_large_en_uncased/3",
     },
 }

keras_hub/src/models/albert/albert_text_classifier.py

@@ -20,10 +20,10 @@ from keras_hub.src.models.text_classifier import TextClassifier
 class AlbertTextClassifier(TextClassifier):
     """An end-to-end ALBERT model for classification tasks
 
-    This model attaches a classification head to a `keras_hub.model.AlbertBackbone`
-    backbone, mapping from the backbone outputs to logit output suitable for
-    a classification task. For usage of this model with pre-trained weights, see
-    the `from_preset()` method.
+    This model attaches a classification head to a
+    `keras_hub.model.AlbertBackbone` backbone, mapping from the backbone outputs
+    to logit output suitable for a classification task. For usage of this model
+    with pre-trained weights, see the `from_preset()` method.
 
     This model can optionally be configured with a `preprocessor` layer, in
     which case it will automatically apply preprocessing to raw inputs during
@@ -36,9 +36,9 @@ class AlbertTextClassifier(TextClassifier):
     Args:
         backbone: A `keras_hub.models.AlertBackbone` instance.
         num_classes: int. Number of classes to predict.
-        preprocessor: A `keras_hub.models.AlbertTextClassifierPreprocessor` or `None`. If
-            `None`, this model will not apply preprocessing, and inputs should
-            be preprocessed before calling the model.
+        preprocessor: A `keras_hub.models.AlbertTextClassifierPreprocessor` 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.

keras_hub/src/models/backbone.py

@@ -1,15 +1,9 @@
-import os
-
 import keras
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.utils.keras_utils import assert_quantization_support
-from keras_hub.src.utils.preset_utils import CONFIG_FILE
-from keras_hub.src.utils.preset_utils import MODEL_WEIGHTS_FILE
 from keras_hub.src.utils.preset_utils import builtin_presets
 from keras_hub.src.utils.preset_utils import get_preset_loader
-from keras_hub.src.utils.preset_utils import save_metadata
-from keras_hub.src.utils.preset_utils import save_serialized_object
+from keras_hub.src.utils.preset_utils import get_preset_saver
 from keras_hub.src.utils.python_utils import classproperty
 
 
@@ -88,10 +82,6 @@ class Backbone(keras.Model):
     def token_embedding(self, value):
         self._token_embedding = value
 
-    def quantize(self, mode, **kwargs):
-        assert_quantization_support()
-        return super().quantize(mode, **kwargs)
-
     def get_config(self):
         # Don't chain to super here. `get_config()` for functional models is
         # a nested layer config and cannot be passed to Backbone constructors.
@@ -193,9 +183,8 @@ class Backbone(keras.Model):
         Args:
             preset_dir: The path to the local model preset directory.
         """
-        save_serialized_object(self, preset_dir, config_file=CONFIG_FILE)
-        self.save_weights(os.path.join(preset_dir, MODEL_WEIGHTS_FILE))
-        save_metadata(self, preset_dir)
+        saver = get_preset_saver(preset_dir)
+        saver.save_backbone(self)
 
     def enable_lora(self, rank):
         """Enable Lora on the backbone.

keras_hub/src/models/bart/bart_backbone.py

@@ -22,9 +22,9 @@ class BartBackbone(Backbone):
     described in
     ["BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension"](https://arxiv.org/abs/1910.13461).
 
-    The default constructor gives a fully customizable, randomly initialized BART
-    model with any number of layers, heads, and embedding dimensions. To load
-    preset architectures and weights, use the `from_preset` constructor.
+    The default constructor gives a fully customizable, randomly initialized
+    BART model with any number of layers, heads, and embedding dimensions. To
+    load preset architectures and weights, use the `from_preset` constructor.
 
     Disclaimer: Pre-trained models are provided on an "as is" basis, without
     warranties or conditions of any kind. The underlying model is provided by a
@@ -78,7 +78,7 @@ class BartBackbone(Backbone):
     )
     output = model(input_data)
     ```
-    """
+    """  # noqa: E501
 
     def __init__(
         self,

keras_hub/src/models/bart/bart_presets.py

@@ -8,11 +8,9 @@ backbone_presets = {
                 "Trained on BookCorpus, English Wikipedia and CommonCrawl."
             ),
             "params": 139417344,
-            "official_name": "BART",
             "path": "bart",
-            "model_card": "https://github.com/facebookresearch/fairseq/blob/main/examples/bart/README.md",
         },
-        "kaggle_handle": "kaggle://keras/bart/keras/bart_base_en/2",
+        "kaggle_handle": "kaggle://keras/bart/keras/bart_base_en/3",
     },
     "bart_large_en": {
         "metadata": {
@@ -21,9 +19,7 @@ backbone_presets = {
                 "Trained on BookCorpus, English Wikipedia and CommonCrawl."
             ),
             "params": 406287360,
-            "official_name": "BART",
             "path": "bart",
-            "model_card": "https://github.com/facebookresearch/fairseq/blob/main/examples/bart/README.md",
         },
         "config": {
             "vocabulary_size": 50265,
@@ -34,7 +30,7 @@ backbone_presets = {
             "dropout": 0.1,
             "max_sequence_length": 1024,
         },
-        "kaggle_handle": "kaggle://keras/bart/keras/bart_large_en/2",
+        "kaggle_handle": "kaggle://keras/bart/keras/bart_large_en/3",
     },
     "bart_large_en_cnn": {
         "metadata": {
@@ -43,9 +39,7 @@ backbone_presets = {
                 "summarization dataset."
             ),
             "params": 406287360,
-            "official_name": "BART",
             "path": "bart",
-            "model_card": "https://github.com/facebookresearch/fairseq/blob/main/examples/bart/README.md",
         },
         "config": {
             "vocabulary_size": 50264,
@@ -56,6 +50,6 @@ backbone_presets = {
             "dropout": 0.1,
             "max_sequence_length": 1024,
         },
-        "kaggle_handle": "kaggle://keras/bart/keras/bart_large_en_cnn/2",
+        "kaggle_handle": "kaggle://keras/bart/keras/bart_large_en_cnn/3",
     },
 }

keras_hub/src/models/bart/bart_seq_2_seq_lm.py

@@ -60,7 +60,8 @@ class BartSeq2SeqLM(Seq2SeqLM):
     bart_lm.generate("The quick brown fox", max_length=30)
     ```
 
-    Use `generate()` with encoder inputs and an incomplete decoder input (prompt).
+    Use `generate()` with encoder inputs and an incomplete decoder input
+    (prompt).
     ```python
     bart_lm = keras_hub.models.BartSeq2SeqLM.from_preset("bart_base_en")
     bart_lm.generate(
@@ -79,10 +80,10 @@ class BartSeq2SeqLM(Seq2SeqLM):
     prompt = {
         "encoder_token_ids": np.array([[0, 133, 2119, 6219, 23602, 2, 1, 1]]),
         "encoder_padding_mask": np.array(
-            [[True, True, True, True, True, True, False, False]]
+            [[1, 1, 1, 1, 1, 1, 0, 0]]
         ),
         "decoder_token_ids": np.array([[2, 0, 133, 1769, 2, 1, 1]]),
-        "decoder_padding_mask": np.array([[True, True, True, True, False, False]])
+        "decoder_padding_mask": np.array([[1, 1, 1, 1, 0, 0]])
     }
 
     bart_lm = keras_hub.models.BartSeq2SeqLM.from_preset(
@@ -95,7 +96,7 @@ class BartSeq2SeqLM(Seq2SeqLM):
     Call `fit()` on a single batch.
     ```python
     features = {
-        "encoder_text": ["The quick brown fox jumped.", "I forgot my homework."],
+        "encoder_text": ["The quick fox jumped.", "I forgot my homework."],
         "decoder_text": ["The fast hazel fox leapt.", "I forgot my assignment."]
     }
     bart_lm = keras_hub.models.BartSeq2SeqLM.from_preset("bart_base_en")
@@ -195,7 +196,7 @@ class BartSeq2SeqLM(Seq2SeqLM):
         cross_attention_cache=None,
         cross_attention_cache_update_index=None,
     ):
-        """Forward pass with a key/value caches for generative decoding..
+        """Forward pass with a key/value caches for generative decoding.
 
         `call_decoder_with_cache` adds an additional inference-time forward pass
         for the model for seq2seq text generation. Unlike calling the model
@@ -241,7 +242,7 @@ class BartSeq2SeqLM(Seq2SeqLM):
             key/value cache in the decoder's self-attention layer and
             `cross_attention_cache` is the key/value cache in the decoder's
             cross-attention layer.
-        """
+        """  # noqa: E501
         # Embedding layers.
         tokens = self.backbone.token_embedding(decoder_token_ids)
         positions = self.backbone.decoder_position_embedding(
@@ -331,7 +332,7 @@ class BartSeq2SeqLM(Seq2SeqLM):
     def _build_cache(
         self, encoder_token_ids, encoder_padding_mask, decoder_token_ids
     ):
-        """Builds the self-attention cache and the cross-attention cache (key/value pairs)."""
+        """Builds the self-attention cache and the cross-attention cache."""
         encoder_hidden_states = self.call_encoder(
             token_ids=encoder_token_ids, padding_mask=encoder_padding_mask
         )
@@ -417,7 +418,7 @@ class BartSeq2SeqLM(Seq2SeqLM):
             prompt = ops.slice(prompt, [0, cache_index], [num_samples, 1])
 
             def repeat_tensor(x):
-                """Repeats tensors along batch axis to match dim for beam search."""
+                """Repeats along batch axis to match dim for beam search."""
                 if ops.shape(x)[0] == num_samples:
                     return x
                 return ops.repeat(x, repeats=num_samples // batch_size, axis=0)

keras_hub/src/models/basnet/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.basnet.basnet_backbone import BASNetBackbone
+from keras_hub.src.models.basnet.basnet_presets import basnet_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(basnet_presets, BASNetBackbone)

keras_hub/src/models/basnet/basnet.py

@@ -0,0 +1,122 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.basnet.basnet_backbone import BASNetBackbone
+from keras_hub.src.models.basnet.basnet_preprocessor import BASNetPreprocessor
+from keras_hub.src.models.image_segmenter import ImageSegmenter
+
+
+@keras_hub_export("keras_hub.models.BASNetImageSegmenter")
+class BASNetImageSegmenter(ImageSegmenter):
+    """BASNet image segmentation task.
+
+    Args:
+        backbone: A `keras_hub.models.BASNetBackbone` instance.
+        preprocessor: `None`, a `keras_hub.models.Preprocessor` instance,
+            a `keras.Layer` instance, or a callable. If `None` no preprocessing
+            will be applied to the inputs.
+
+    Example:
+    ```python
+    import keras_hub
+
+    images = np.ones(shape=(1, 288, 288, 3))
+    labels = np.zeros(shape=(1, 288, 288, 1))
+
+    image_encoder = keras_hub.models.ResNetBackbone.from_preset(
+        "resnet_18_imagenet",
+        load_weights=False
+    )
+    backbone = keras_hub.models.BASNetBackbone(
+        image_encoder,
+        num_classes=1,
+        image_shape=[288, 288, 3]
+    )
+    model = keras_hub.models.BASNetImageSegmenter(backbone)
+
+    # Evaluate the model
+    pred_labels = model(images)
+
+    # Train the model
+    model.compile(
+        optimizer="adam",
+        loss=keras.losses.BinaryCrossentropy(from_logits=False),
+        metrics=["accuracy"],
+    )
+    model.fit(images, labels, epochs=3)
+    ```
+    """
+
+    backbone_cls = BASNetBackbone
+    preprocessor_cls = BASNetPreprocessor
+
+    def __init__(
+        self,
+        backbone,
+        preprocessor=None,
+        **kwargs,
+    ):
+        # === Functional Model ===
+        x = backbone.input
+        outputs = backbone(x)
+        # only return the refinement module's output as final prediction
+        outputs = outputs["refine_out"]
+        super().__init__(inputs=x, outputs=outputs, **kwargs)
+
+        # === Config ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+
+    def compute_loss(self, x, y, y_pred, *args, **kwargs):
+        # train BASNet's prediction and refinement module outputs against the
+        # same ground truth data
+        outputs = self.backbone(x)
+        losses = []
+        for output in outputs.values():
+            losses.append(super().compute_loss(x, y, output, *args, **kwargs))
+        return keras.ops.sum(losses, axis=0)
+
+    def compile(
+        self,
+        optimizer="auto",
+        loss="auto",
+        metrics="auto",
+        **kwargs,
+    ):
+        """Configures the `BASNet` task for training.
+
+        `BASNet` extends the default compilation signature
+        of `keras.Model.compile` with defaults for `optimizer` and `loss`. To
+        override these defaults, pass any value to these arguments during
+        compilation.
+
+        Args:
+            optimizer: `"auto"`, an optimizer name, or a `keras.Optimizer`
+                instance. Defaults to `"auto"`, which uses the default
+                optimizer for `BASNet`. See `keras.Model.compile` and
+                `keras.optimizers` for more info on possible `optimizer`
+                values.
+            loss: `"auto"`, a loss name, or a `keras.losses.Loss` instance.
+                Defaults to `"auto"`, in which case the default loss
+                computation of `BASNet` will be applied.
+                See `keras.Model.compile` and `keras.losses` for more info on
+                possible `loss` values.
+            metrics: `"auto"`, or a list of metrics to be evaluated by
+                the model during training and testing. Defaults to `"auto"`,
+                where a `keras.metrics.Accuracy` will be applied to track the
+                accuracy of the model during training.
+                See `keras.Model.compile` and `keras.metrics` for
+                more info on possible `metrics` values.
+            **kwargs: See `keras.Model.compile` for a full list of arguments
+                supported by the compile method.
+        """
+        if loss == "auto":
+            loss = keras.losses.BinaryCrossentropy()
+        if metrics == "auto":
+            metrics = [keras.metrics.Accuracy()]
+        super().compile(
+            optimizer=optimizer,
+            loss=loss,
+            metrics=metrics,
+            **kwargs,
+        )