keras-hub-nightly 0.16.1.dev202409230338__py3-none-any.whl → 0.16.1.dev202409250340__py3-none-any.whl

keras_hub/api/layers/__init__.py

@@ -56,6 +56,8 @@ from keras_hub.src.models.pali_gemma.pali_gemma_image_converter import (
 from keras_hub.src.models.resnet.resnet_image_converter import (
     ResNetImageConverter,
 )
+from keras_hub.src.models.sam.sam_mask_decoder import SAMMaskDecoder
+from keras_hub.src.models.sam.sam_prompt_encoder import SAMPromptEncoder
 from keras_hub.src.models.whisper.whisper_audio_converter import (
     WhisperAudioConverter,
 )

keras_hub/api/models/__init__.py

@@ -175,6 +175,7 @@ from keras_hub.src.models.image_classifier import ImageClassifier
 from keras_hub.src.models.image_classifier_preprocessor import (
     ImageClassifierPreprocessor,
 )
+from keras_hub.src.models.image_segmenter import ImageSegmenter
 from keras_hub.src.models.llama3.llama3_backbone import Llama3Backbone
 from keras_hub.src.models.llama3.llama3_causal_lm import Llama3CausalLM
 from keras_hub.src.models.llama3.llama3_causal_lm_preprocessor import (
@@ -255,6 +256,8 @@ from keras_hub.src.models.roberta.roberta_text_classifier_preprocessor import (
     RobertaTextClassifierPreprocessor as RobertaPreprocessor,
 )
 from keras_hub.src.models.roberta.roberta_tokenizer import RobertaTokenizer
+from keras_hub.src.models.sam.sam_backbone import SAMBackbone
+from keras_hub.src.models.sam.sam_image_segmenter import SAMImageSegmenter
 from keras_hub.src.models.seq_2_seq_lm import Seq2SeqLM
 from keras_hub.src.models.seq_2_seq_lm_preprocessor import Seq2SeqLMPreprocessor
 from keras_hub.src.models.t5.t5_backbone import T5Backbone

keras_hub/src/models/image_segmenter.py

@@ -0,0 +1,86 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.task import Task
+
+
+@keras_hub_export("keras_hub.models.ImageSegmenter")
+class ImageSegmenter(Task):
+    """Base class for all image segmentation tasks.
+
+    `ImageSegmenter` tasks wrap a `keras_hub.models.Task` and
+    a `keras_hub.models.Preprocessor` to create a model that can be used for
+    image segmentation.
+
+    All `ImageSegmenter` tasks include a `from_preset()` constructor which can
+    be used to load a pre-trained config and weights.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Default compilation.
+        self.compile()
+
+    def compile(
+        self,
+        optimizer="auto",
+        loss="auto",
+        *,
+        metrics="auto",
+        **kwargs,
+    ):
+        """Configures the `ImageSegmenter` task for training.
+
+        The `ImageSegmenter` task extends the default compilation signature of
+        `keras.Model.compile` with defaults for `optimizer`, `loss`, and
+        `metrics`. 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 the given model and task. 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"`, where a
+                `keras.losses.SparseCategoricalCrossentropy` loss will be
+                applied for the classification task. 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.SparseCategoricalAccuracy` 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 optimizer == "auto":
+            optimizer = keras.optimizers.Adam(5e-5)
+        if loss == "auto":
+            activation = getattr(self, "activation", None)
+            activation = keras.activations.get(activation)
+            from_logits = activation != keras.activations.softmax
+            loss = keras.losses.CategoricalCrossentropy(from_logits=from_logits)
+        if metrics == "auto":
+            metrics = [keras.metrics.CategoricalAccuracy()]
+        super().compile(
+            optimizer=optimizer,
+            loss=loss,
+            metrics=metrics,
+            **kwargs,
+        )

keras_hub/src/models/sam/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

keras_hub/src/models/sam/sam_backbone.py

@@ -0,0 +1,153 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+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.SAMBackbone")
+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.
+        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
+            generate segmentation masks given the embeddings generated by the
+            backbone and the prompt encoder.
+        dtype: The dtype of the layer weights.
+
+    Example:
+    ```python
+    image_size=128
+    batch_size=2
+    input_data = {
+        "images": np.ones(
+            (batch_size, image_size, image_size, 3),
+            dtype="float32",
+        ),
+        "points": np.ones((batch_size, 1, 2), dtype="float32"),
+        "labels": np.ones((batch_size, 1), dtype="float32"),
+        "boxes": np.ones((batch_size, 1, 2, 2), dtype="float32"),
+        "masks": np.zeros(
+            (batch_size, 0, image_size, image_size, 1)
+        ),
+    }
+    image_encoder = keras_hub.models.ViTDetBackbone(
+        hidden_size=16,
+        num_layers=16,
+        intermediate_dim=16 * 4,
+        num_heads=16,
+        global_attention_layer_indices=[2, 5, 8, 11],
+        patch_size=16,
+        num_output_channels=8,
+        window_size=2,
+        image_shape=(image_size, image_size, 3),
+    )
+    prompt_encoder = keras_hub.layers.SAMPromptEncoder(
+        hidden_size=8,
+        image_embedding_size=(8, 8),
+        input_image_size=(
+            image_size,
+            image_size,
+        ),
+        mask_in_channels=16,
+    )
+    mask_decoder = keras_hub.layers.SAMMaskDecoder(
+        num_layers=2,
+        hidden_size=8,
+        intermediate_dim=32,
+        num_heads=8,
+        embedding_dim=8,
+        num_multimask_outputs=3,
+        iou_head_depth=3,
+        iou_head_hidden_dim=8,
+    )
+    backbone = keras_hub.models.SAMBackbone(
+        image_encoder=image_encoder,
+        prompt_encoder=prompt_encoder,
+        mask_decoder=mask_decoder,
+        image_shape=(image_size, image_size, 3),
+    )
+    backbone(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        image_encoder,
+        prompt_encoder,
+        mask_decoder,
+        dtype=None,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.image_encoder = image_encoder
+        self.prompt_encoder = prompt_encoder
+        self.mask_decoder = mask_decoder
+        # === Functional model
+        image_input = self.image_encoder.input
+
+        inputs = {
+            "images": image_input,
+            "points": keras.Input(shape=[None, 2], name="points"),
+            "labels": keras.Input(shape=[None], name="labels"),
+            "boxes": keras.Input(shape=[None, 2, 2], name="boxes"),
+            "masks": keras.Input(shape=[None, None, None, 1], name="masks"),
+        }
+        image_embeddings = self.image_encoder.output
+        prompt_embeddings = self.prompt_encoder(**inputs)
+        outputs = {
+            "image_embeddings": image_embeddings,
+        }
+        outputs.update(prompt_embeddings)
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            dtype=dtype,
+            **kwargs,
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "image_encoder": keras.layers.serialize(self.image_encoder),
+                "prompt_encoder": keras.layers.serialize(self.prompt_encoder),
+                "mask_decoder": keras.layers.serialize(self.mask_decoder),
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        config.update(
+            {
+                "image_encoder": keras.layers.deserialize(
+                    config["image_encoder"]
+                ),
+                "prompt_encoder": keras.layers.deserialize(
+                    config["prompt_encoder"]
+                ),
+                "mask_decoder": keras.layers.deserialize(
+                    config["mask_decoder"]
+                ),
+            }
+        )
+
+        return super().from_config(config)

keras_hub/src/models/sam/sam_image_segmenter.py

@@ -0,0 +1,237 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.image_segmenter import ImageSegmenter
+from keras_hub.src.models.sam.sam_backbone import SAMBackbone
+
+
+@keras_hub_export("keras_hub.models.SAMImageSegmenter")
+class SAMImageSegmenter(ImageSegmenter):
+    """The Segment Anything (SAM) image segmenter Model.
+
+    SAM works by prompting the input images. There are three ways to prompt:
+    (1) Labelled Points: Foreground points (points with label 1) are encoded
+        such that the output masks generated by the mask decoder contain them
+        and background points (points with label 0) are encoded such that the
+        generated masks don't contain them.
+    (2) Box: A box tells the model which part/crop of the image to segment.
+    (3) Mask: An input mask can be used to refine the output of the mask
+        decoder.
+    These prompts can be mixed and matched but at least one of the prompts
+    must be present. To turn off a particular prompt, simply exclude it from
+    the inputs to the model.
+    (1) For points prompts, the expected shape is `(batch, num_points, 2)`.
+        The labels must have a corresponding shape of `(batch, num_points)`.
+    (2) For box prompt, the expected shape is `(batch, 1, 2, 2)`.
+    (3) Similarly, mask prompts have shape `(batch, 1, H, W, 1)`.
+
+
+    Args:
+      backbone: A `keras_hub.models.VGGBackbone` instance.
+
+    Example:
+    Load pretrained model using `from_preset`.
+
+    ```python
+    image_size=128
+    batch_size=2
+    input_data = {
+        "images": np.ones(
+            (batch_size, image_size, image_size, 3),
+            dtype="float32",
+        ),
+        "points": np.ones((batch_size, 1, 2), dtype="float32"),
+        "labels": np.ones((batch_size, 1), dtype="float32"),
+        "boxes": np.ones((batch_size, 1, 2, 2), dtype="float32"),
+        "masks": np.zeros(
+            (batch_size, 0, image_size, image_size, 1)
+        ),
+    }
+    # todo: update preset name
+    sam = keras_hub.models.SAMImageSegmenter.from_preset(`sam_base`)
+    sam(input_data)
+    ```
+
+    Load segment anything image segmenter with custom backbone
+
+    ```python
+    image_size = 128
+    batch_size = 2
+    images = np.ones(
+        (batch_size, image_size, image_size, 3),
+        dtype="float32",
+    )
+    image_encoder = ViTDetBackbone(
+        hidden_size=16,
+        num_layers=16,
+        intermediate_dim=16 * 4,
+        num_heads=16,
+        global_attention_layer_indices=[2, 5, 8, 11],
+        patch_size=16,
+        num_output_channels=8,
+        window_size=2,
+        image_shape=(image_size, image_size, 3),
+    )
+    prompt_encoder = SAMPromptEncoder(
+        hidden_size=8,
+        image_embedding_size=(8, 8),
+        input_image_size=(
+            image_size,
+            image_size,
+        ),
+        mask_in_channels=16,
+    )
+    mask_decoder = SAMMaskDecoder(
+        num_layers=2,
+        hidden_size=8,
+        intermediate_dim=32,
+        num_heads=8,
+        embedding_dim=8,
+        num_multimask_outputs=3,
+        iou_head_depth=3,
+        iou_head_hidden_dim=8,
+    )
+    backbone = SAMBackbone(
+        image_encoder=image_encoder,
+        prompt_encoder=prompt_encoder,
+        mask_decoder=mask_decoder,
+        image_shape=(image_size, image_size, 3),
+    )
+    sam = SAMImageSegmenter(
+        backbone=backbone
+    )
+    ```
+
+    For example, to pass in all the prompts, do:
+
+    ```python
+
+    points = np.array([[[512., 512.], [100., 100.]]])
+    # For labels: 1 means foreground point, 0 means background
+    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:
+    inputs = {
+        "images": image,
+        "points": points,
+        "labels": labels,
+        "boxes": box,
+        "masks": input_mask
+    }
+    outputs = sam.predict(inputs)
+    masks, iou_pred = outputs["masks"], outputs["iou_pred"]
+    ```
+
+    The first mask in the output `masks` (i.e. `masks[:, 0, ...]`) is the best
+    mask predicted by the model based on the prompts. Other `masks`
+    (i.e. `masks[:, 1:, ...]`) are alternate predictions that can be used if
+    they are desired over the first one.
+    Now, in case of only points and box prompts, simply exclude the masks:
+
+    ```python
+    inputs = {
+        "images": image,
+        "points": points,
+        "labels": labels,
+        "boxes": box,
+    }
+
+    outputs = sam.predict(inputs)
+    masks, iou_pred = outputs["masks"], outputs["iou_pred"]
+    ```
+
+    Another example is that only points prompts are present.
+    Note that if point prompts are present but no box prompt is present, the
+    points must be padded using a zero point and -1 label:
+
+    ```python
+    padded_points = np.concatenate(
+        [points, np.zeros((1, 1, 2))], axis=1
+    )
+
+    padded_labels = np.concatenate(
+        [labels, -np.ones((1, 1))], axis=1
+    )
+    inputs = {
+        "images": image,
+        "points": padded_points,
+        "labels": padded_labels,
+    }
+    outputs = sam.predict(inputs)
+    masks, iou_pred = outputs["masks"], outputs["iou_pred"]
+    ```
+    """
+
+    backbone_cls = SAMBackbone
+    preprocessor_cls = None
+
+    def __init__(self, backbone, preprocessor=None, **kwargs):
+        # The implementation has been adapted form [Segment Anything
+        # paper](https://arxiv.org/abs/2304.02643) and [Segment Anything
+        # GitHub](https://github.com/facebookresearch/segment-anything) and
+        # [Detectron2](https://github.com/facebookresearch/detectron2).
+        # === Layers ===
+        self.backbone = backbone
+        # === Functional Model ===
+        inputs = self.backbone.input
+        x = self.backbone(inputs)
+        outputs = self.backbone.mask_decoder(**x)
+        super().__init__(inputs=inputs, outputs=outputs, **kwargs)
+
+    def predict_step(self, *args, **kwargs):
+        if len(args) == 2:
+            args = (args[0], self._add_placeholder_prompts(args[-1]))
+        else:
+            args = (self._add_placeholder_prompts(args[0]),)
+
+        return super().predict_step(*args, **kwargs)
+
+    def fit(self, *args, **kwargs):
+        raise NotImplementedError(
+            "Segment Anything Model only supports inference for now. Training"
+            " the model isn't supported yet."
+        )
+
+    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.
+        """
+        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).
+        zeros = ops.zeros if ops.is_tensor(inputs["images"]) else np.zeros
+
+        # Fill in missing inputs.
+        if "points" not in inputs:
+            inputs["points"] = zeros((batch_size, 0, 2))
+        if "labels" not in inputs:
+            inputs["labels"] = zeros((batch_size, 0))
+        if "boxes" not in inputs:
+            inputs["boxes"] = zeros((batch_size, 0, 2, 2))
+        if "masks" not in inputs:
+            inputs["masks"] = zeros((batch_size, 0, 256, 256, 1))
+
+        return inputs