keras-hub-nightly 0.16.1.dev202409240339__py3-none-any.whl → 0.16.1.dev202409260340__py3-none-any.whl

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

keras_hub/src/models/sam/sam_layers.py

@@ -0,0 +1,402 @@
+# 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 math
+
+import keras
+from keras import ops
+
+
+class MLP(keras.layers.Layer):
+    """A MLP block with architecture.
+
+    `input_dim -> [hidden_dim] * (num_layers - 1) -> output_dim`.
+
+    Args:
+        hidden_dim: int. The number of units in the hidden layers.
+        output_dim: int. The number of units in the output layer.
+        num_layers: int. The total number of dense layers to use.
+        activation: str. Activation to use in the hidden layers.
+            Default is `"relu"`.
+    """
+
+    def __init__(
+        self, hidden_dim, output_dim, num_layers, activation="relu", **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.num_layers = num_layers
+        self.activation = activation
+        h = [hidden_dim] * (num_layers - 1)
+        self.mlp_block = []
+        for hidden_dim in h:
+            self.mlp_block.append(
+                keras.layers.Dense(hidden_dim, dtype=self.dtype_policy)
+            )
+            self.mlp_block.append(
+                keras.layers.Activation(activation, dtype=self.dtype_policy)
+            )
+        self.mlp_block.append(
+            keras.layers.Dense(output_dim, dtype=self.dtype_policy)
+        )
+        self.mlp_block = keras.models.Sequential(self.mlp_block)
+
+    def build(self, input_shape):
+        self.mlp_block.build(input_shape)
+        self.built = True
+
+    def call(self, x):
+        return self.mlp_block(x)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "output_dim": self.output_dim,
+                "num_layers": self.num_layers,
+                "activation": self.activation,
+            }
+        )
+        return config
+
+
+class MultiHeadAttentionWithDownsampling(keras.layers.Layer):
+    """Multi-Head Attention with downsampling.
+
+    An attention layer that allows for downscaling the size of the embedding
+    after projection to queries, keys, and values.
+    This layer first downscales the features of input queries, keys, and
+    values using a dense layer. Multi-head attention is then performed
+    and the attention map is projected back (upscaled) to the number of
+    input features.
+
+    Args:
+        num_heads: int. Number of attention heads.
+        key_dim: int. Size of each attention head for query, key, and
+            value.
+        downsample_rate: int, optional. The factor by which to downscale the
+            input features i.e. the input features of size `key_dim` are
+            projected down to `key_dim // downsample_rate`.
+    """
+
+    def __init__(self, num_heads, key_dim, downsample_rate=1, **kwargs):
+        super().__init__(**kwargs)
+        self.num_heads = num_heads
+        self.key_dim = key_dim
+        self.downsample_rate = downsample_rate
+        self.internal_dims = key_dim // downsample_rate
+
+        # Downsample
+        self.query_proj = keras.layers.Dense(
+            self.internal_dims * self.num_heads, dtype=self.dtype_policy
+        )
+        self.key_proj = keras.layers.Dense(
+            self.internal_dims * self.num_heads, dtype=self.dtype_policy
+        )
+        self.value_proj = keras.layers.Dense(
+            self.internal_dims * self.num_heads, dtype=self.dtype_policy
+        )
+
+        # Upsample
+        self.out_proj = keras.layers.Dense(
+            self.key_dim * self.num_heads, dtype=self.dtype_policy
+        )
+
+    def build(self, input_shape=None):
+        self.query_proj.build([None, None, self.num_heads * self.key_dim])
+        self.key_proj.build([None, None, self.num_heads * self.key_dim])
+        self.value_proj.build([None, None, self.num_heads * self.key_dim])
+        self.out_proj.build([None, None, self.internal_dims * self.num_heads])
+        self.built = True
+
+    def _separate_heads(self, x):
+        shape = ops.shape(x)
+        batch_size, N, channels = shape[0], shape[1], shape[2]
+        x = ops.reshape(
+            x, (batch_size, N, self.num_heads, channels // self.num_heads)
+        )
+        return ops.transpose(x, axes=(0, 2, 1, 3))
+
+    def _recombine_heads(self, x):
+        shape = ops.shape(x)
+        batch_size, num_heads, N_T, channels_per_head = (
+            shape[0],
+            shape[1],
+            shape[2],
+            shape[3],
+        )
+        x = ops.transpose(x, axes=(0, 2, 1, 3))
+        return ops.reshape(x, (batch_size, N_T, num_heads * channels_per_head))
+
+    def call(self, query, value, key):
+        query = self.query_proj(query)
+        key = self.key_proj(key)
+        value = self.value_proj(value)
+
+        # Separate into heads
+        query = self._separate_heads(query)
+        key = self._separate_heads(key)
+        value = self._separate_heads(value)
+
+        # Attention
+        channels_per_head = ops.shape(query)[-1]
+        out = ops.matmul(query, ops.transpose(key, (0, 1, 3, 2)))
+        out = out / ops.sqrt(
+            ops.cast(channels_per_head, dtype=self.compute_dtype)
+        )
+        out = ops.softmax(out, axis=-1)
+
+        # Get output
+        attention_map = out @ value
+        attention_map = self._recombine_heads(attention_map)
+        return self.out_proj(attention_map)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_heads": self.num_heads,
+                "key_dim": self.key_dim,
+                "downsample_rate": self.downsample_rate,
+            }
+        )
+        return config
+
+
+class TwoWayMultiHeadAttention(keras.layers.Layer):
+    """Two-way multi-head attention layer.
+
+    Args:
+        num_heads: int. Number of attention heads.
+        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.
+        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
+            layer. Defaults to "relu".
+    """
+
+    def __init__(
+        self,
+        num_heads,
+        key_dim,
+        intermediate_dim,
+        skip_first_layer_pos_embedding,
+        attention_downsample_rate=2,
+        activation="relu",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_heads = num_heads
+        self.key_dim = key_dim
+        self.intermediate_dim = intermediate_dim
+        self.skip_first_layer_pos_embedding = skip_first_layer_pos_embedding
+        self.attention_downsample_rate = attention_downsample_rate
+        self.activation = activation
+
+        self.self_attention = MultiHeadAttentionWithDownsampling(
+            num_heads=num_heads, key_dim=key_dim, dtype=self.dtype_policy
+        )
+        self.layer_norm1 = keras.layers.LayerNormalization(
+            epsilon=1e-5, dtype=self.dtype_policy
+        )
+        self.cross_attention_token_to_image = (
+            MultiHeadAttentionWithDownsampling(
+                num_heads=num_heads,
+                key_dim=key_dim,
+                downsample_rate=attention_downsample_rate,
+                dtype=self.dtype_policy,
+            )
+        )
+        self.layer_norm2 = keras.layers.LayerNormalization(
+            epsilon=1e-5, dtype=self.dtype_policy
+        )
+
+        self.mlp_block = MLP(
+            intermediate_dim,
+            key_dim * num_heads,
+            num_layers=2,
+            activation=activation,
+            dtype=self.dtype_policy,
+        )
+
+        self.layer_norm3 = keras.layers.LayerNormalization(
+            epsilon=1e-5, dtype=self.dtype_policy
+        )
+        self.cross_attention_image_to_token = (
+            MultiHeadAttentionWithDownsampling(
+                num_heads=num_heads,
+                key_dim=key_dim,
+                downsample_rate=attention_downsample_rate,
+                dtype=self.dtype_policy,
+            )
+        )
+        self.layer_norm4 = keras.layers.LayerNormalization(
+            epsilon=1e-5, dtype=self.dtype_policy
+        )
+
+    def build(self, input_shape=None):
+        self.self_attention.build()
+        self.layer_norm1.build([None, None, self.num_heads * self.key_dim])
+        self.cross_attention_token_to_image.build()
+        self.layer_norm2.build([None, None, self.num_heads * self.key_dim])
+        self.mlp_block.build([None, None, self.num_heads * self.key_dim])
+        self.layer_norm3.build([None, None, self.num_heads * self.key_dim])
+        self.cross_attention_image_to_token.build()
+        self.layer_norm4.build([None, None, self.num_heads * self.key_dim])
+        self.built = True
+
+    def call(self, queries, keys, query_pos_embedding, key_pos_embedding):
+        if self.skip_first_layer_pos_embedding:
+            queries = self.self_attention(
+                query=queries, value=queries, key=queries
+            )
+        else:
+            queries_with_pos_embedding = queries + query_pos_embedding
+            attention_map = self.self_attention(
+                query=queries_with_pos_embedding,
+                key=queries_with_pos_embedding,
+                value=queries,
+            )
+            queries = queries + attention_map
+        queries = self.layer_norm1(queries)
+
+        queries_with_pos_embedding = queries + query_pos_embedding
+        keys_with_pos_embedding = keys + key_pos_embedding
+        attention_map = self.cross_attention_token_to_image(
+            query=queries_with_pos_embedding,
+            key=keys_with_pos_embedding,
+            value=keys,
+        )
+        queries = queries + attention_map
+        queries = self.layer_norm2(queries)
+
+        mlp_out = self.mlp_block(queries)
+        queries = queries + mlp_out
+        queries = self.layer_norm3(queries)
+
+        queries_with_pos_embedding = queries + query_pos_embedding
+        keys_with_pos_embedding = keys + key_pos_embedding
+        attention_map = self.cross_attention_image_to_token(
+            query=keys_with_pos_embedding,
+            key=queries_with_pos_embedding,
+            value=queries,
+        )
+        keys = keys + attention_map
+        keys = self.layer_norm4(keys)
+
+        return queries, keys
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "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,
+                "attention_downsample_rate": self.attention_downsample_rate,
+                "activation": self.activation,
+            }
+        )
+        return config
+
+
+class RandomFrequencyPositionalEmbeddings(keras.layers.Layer):
+    """Positional encoding using random spatial frequencies.
+
+    This layer maps coordinates/points in 2D space to positional
+    encodings using random spatial frequencies.
+
+    Args:
+        num_positional_features: int. Number of positional features
+            in the output.
+        scale: float. The standard deviation of the random frequencies.
+    """
+
+    def __init__(self, num_positional_features, scale, **kwargs):
+        super().__init__(**kwargs)
+        self.num_positional_features = num_positional_features
+        self.scale = scale
+        self.positional_encoding_gaussian_matrix = self.add_weight(
+            name="positional_encoding_gaussian_matrix",
+            shape=(2, self.num_positional_features),
+            dtype=self.variable_dtype,
+            trainable=False,
+            initializer=keras.initializers.get("normal"),
+        )
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def _positional_encodings(self, coords):
+        coords = coords * 2 - 1
+        coords = coords @ ops.cast(
+            self.positional_encoding_gaussian_matrix, dtype=self.compute_dtype
+        )
+        coords = coords * (2 * math.pi)
+        return ops.concatenate([ops.sin(coords), ops.cos(coords)], axis=-1)
+
+    def call(self, size):
+        return self.encode_image(size)
+
+    def encode_image(self, size):
+        """Generate a positional encoding for an image of any given size.
+        Args:
+            size: tuple[int, int]. The size of the image.
+        Returns:
+            tensor: Positional encoding of the image.
+        """
+        height, width = size
+        grid = ops.ones(shape=(height, width), dtype=self.compute_dtype)
+        y_embed = ops.cumsum(grid, axis=0) - 0.5
+        x_embed = ops.cumsum(grid, axis=1) - 0.5
+        y_embed = y_embed / ops.cast(height, self.compute_dtype)
+        x_embed = x_embed / ops.cast(width, self.compute_dtype)
+        return self._positional_encodings(
+            ops.stack([x_embed, y_embed], axis=-1)
+        )
+
+    def encode_coordinates(self, coords_input, image_size):
+        """Positionally encode points that are not normalized to `[0, 1]`.
+        Args:
+            coords_input: tensor. 2D coordinates/points to map.
+            image_size: tuple[int, int]. Height and width of the image
+                being prompted.
+        Returns:
+            tensor: Positional encodings of the normalized coordinates.
+        """
+        coords_normalized = ops.stack(
+            [
+                coords_input[..., 0] / image_size[1],
+                coords_input[..., 1] / image_size[0],
+            ],
+            axis=-1,
+        )
+        return self._positional_encodings(coords_normalized)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_positional_features": self.num_positional_features,
+                "scale": self.scale,
+            }
+        )
+        return config