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_mask_decoder.py

@@ -0,0 +1,270 @@
+# 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 import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.sam.sam_layers import MLP
+from keras_hub.src.models.sam.sam_transformer import TwoWayTransformer
+
+
+@keras_hub_export("keras_hub.layers.SAMMaskDecoder")
+class SAMMaskDecoder(keras.layers.Layer):
+    """Mask decoder for the Segment Anything Model (SAM).
+
+    This lightweight module efficiently maps the image embedding and a set of
+    prompt embeddings to an output mask. Before applying the transformer
+    decoder, the layer first inserts into the set of prompt embeddings a
+    learned output token embedding that will be used at the decoder's output.
+    For simplicity, these embeddings (not including the image embedding) are
+    collectively called "tokens".
+
+    The image embeddings, positional image embeddings, and tokens are passed
+    through a transformer decoder. After running the decoder, the layer
+    upsamples the updated image embedding by 4x with two transposed
+    convolutional layers (now it's downscaled 4x relative to the input
+    image). Then, the tokens attend once more to the image embedding and
+    the updated output token embedding are passed to a small 3-layer MLP that
+    outputs a vector matching the channel dimension of the upscaled image
+    embedding.
+
+    Finally, a mask is predicted with a spatially point-wise
+    product between the upscaled image embedding and the MLP's output.
+
+    Args:
+        hidden_size: int. The hidden size of the TwoWayTransformer.
+        num_layers: int. The number of layers in the TwoWayTransformer.
+        intermediate_dim: int. The intermediate dimension of the
+            TwoWayTransformer.
+        num_heads: int. The number of heads in the TwoWayTransformer.
+        embedding_dim: int, optional. The number of input features to the
+            transformer decoder. Defaults to `256`.
+        num_multimask_outputs: int, optional. Number of multimask outputs.
+            The model would generate these many extra masks. The total masks
+            generated by the model are `1 + num_multimask_outputs`. Defaults
+            to `3`.
+        iou_head_depth: int, optional. The depth of the dense net used to
+            predict the IoU confidence score. Defaults to `3`.
+        iou_head_hidden_dim: int, optional. The number of units in the hidden
+            layers used in the dense net to predict the IoU confidence score.
+            Defaults to `256`.
+        activation: str, optional. Activation to use in the mask upscaler
+            network. Defaults to `"gelu"`.
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden_size,
+        num_layers,
+        intermediate_dim,
+        num_heads,
+        embedding_dim=256,
+        num_multimask_outputs=3,
+        iou_head_depth=3,
+        iou_head_hidden_dim=256,
+        activation="gelu",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.intermediate_dim = intermediate_dim
+        self.num_heads = num_heads
+        self.embedding_dim = embedding_dim
+        transformer = TwoWayTransformer(
+            num_layers=num_layers,
+            hidden_size=hidden_size,
+            intermediate_dim=intermediate_dim,
+            num_heads=num_heads,
+            dtype=self.dtype_policy,
+        )
+        self.transformer = transformer
+        self.num_multimask_outputs = num_multimask_outputs
+        self.iou_head_depth = iou_head_depth
+        self.iou_head_hidden_dim = iou_head_hidden_dim
+        self.activation = activation
+
+        self.iou_token = keras.layers.Embedding(
+            1, embedding_dim, dtype=self.dtype_policy
+        )
+        self.num_mask_tokens = num_multimask_outputs + 1
+        self.mask_tokens = keras.layers.Embedding(
+            self.num_mask_tokens, embedding_dim, dtype=self.dtype_policy
+        )
+
+        self.output_upscaling = keras.models.Sequential(
+            [
+                keras.layers.Conv2DTranspose(
+                    embedding_dim // 4,
+                    kernel_size=2,
+                    strides=2,
+                    dtype=self.dtype_policy,
+                ),
+                keras.layers.LayerNormalization(
+                    epsilon=1e-6, dtype=self.dtype_policy
+                ),
+                keras.layers.Activation(activation, dtype=self.dtype_policy),
+                keras.layers.Conv2DTranspose(
+                    embedding_dim // 8,
+                    kernel_size=2,
+                    strides=2,
+                    dtype=self.dtype_policy,
+                ),
+                keras.layers.Activation(activation, dtype=self.dtype_policy),
+            ]
+        )
+
+        self.output_hypernetworks_mlps = [
+            MLP(embedding_dim, embedding_dim // 8, 3, dtype=self.dtype_policy)
+            for _ in range(self.num_mask_tokens)
+        ]
+
+        self.iou_prediction_head = MLP(
+            iou_head_hidden_dim,
+            self.num_mask_tokens,
+            iou_head_depth,
+            dtype=self.dtype_policy,
+        )
+
+    def build(self, input_shape=None, **kwargs):
+        self.transformer.build()
+        self.iou_token.build([None])
+        self.mask_tokens.build([None])
+        self.output_upscaling.build([None, None, None, self.embedding_dim])
+        for mlp in self.output_hypernetworks_mlps:
+            mlp.build([None, self.embedding_dim])
+        self.iou_prediction_head.build([None, self.embedding_dim])
+        self.built = True
+
+    def call(
+        self,
+        image_embeddings,
+        prompt_dense_positional_embeddings,
+        prompt_sparse_embeddings,
+        prompt_dense_embeddings,
+    ):
+        masks, iou_pred = self._predict_masks(
+            image_embeddings=image_embeddings,
+            image_positional_embeddings=prompt_dense_positional_embeddings,
+            prompt_sparse_embeddings=prompt_sparse_embeddings,
+            prompt_dense_embeddings=prompt_dense_embeddings,
+        )
+
+        return {"masks": masks, "iou_pred": iou_pred}
+
+    def _predict_masks(
+        self,
+        image_embeddings,
+        image_positional_embeddings,
+        prompt_sparse_embeddings,
+        prompt_dense_embeddings,
+    ):
+        indices_iou = ops.arange(1, dtype="int32")
+        indices_mask = ops.arange(self.num_mask_tokens, dtype="int32")
+
+        output_tokens = ops.concatenate(
+            [self.iou_token(indices_iou), self.mask_tokens(indices_mask)],
+            axis=0,
+        )
+        output_tokens = ops.broadcast_to(
+            output_tokens[None, ...],
+            shape=(
+                ops.shape(prompt_sparse_embeddings)[0],
+                ops.shape(output_tokens)[0],
+                ops.shape(output_tokens)[1],
+            ),
+        )
+        tokens = ops.concatenate(
+            [output_tokens, prompt_sparse_embeddings], axis=1
+        )
+
+        source = ops.broadcast_to(
+            image_embeddings,
+            shape=(
+                ops.shape(tokens)[0],
+                ops.shape(image_embeddings)[1],
+                ops.shape(image_embeddings)[2],
+                ops.shape(image_embeddings)[3],
+            ),
+        )
+        source = source + prompt_dense_embeddings
+        positional_source = ops.broadcast_to(
+            image_positional_embeddings,
+            shape=(
+                ops.shape(tokens)[0],
+                ops.shape(image_embeddings)[1],
+                ops.shape(image_embeddings)[2],
+                ops.shape(image_embeddings)[3],
+            ),
+        )
+        shape = ops.shape(source)
+        batch_dim, height, width, channels = (
+            shape[0],
+            shape[1],
+            shape[2],
+            shape[3],
+        )
+
+        hidden_state, source = self.transformer(
+            source, positional_source, tokens
+        )
+        iou_token_out = hidden_state[:, 0, :]
+        mask_tokens_out = hidden_state[:, 1 : (1 + self.num_mask_tokens), :]
+
+        source = ops.reshape(source, (batch_dim, height, width, channels))
+        upscaled_embeddings = self.output_upscaling(source)
+        hyper_in_list = []
+        for i in range(self.num_mask_tokens):
+            hyper_in_list.append(
+                self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])
+            )
+        hyper_in = ops.stack(hyper_in_list, axis=1)
+        shape = ops.shape(upscaled_embeddings)
+        batch_dim, height, width, channels = (
+            shape[0],
+            shape[1],
+            shape[2],
+            shape[3],
+        )
+        upscaled_embeddings = ops.reshape(
+            ops.transpose(upscaled_embeddings, axes=(0, 3, 1, 2)),
+            (batch_dim, channels, height * width),
+        )
+        masks = ops.reshape(
+            hyper_in @ upscaled_embeddings,
+            (batch_dim, self.num_mask_tokens, height, width),
+        )
+
+        iou_pred = self.iou_prediction_head(iou_token_out)
+
+        return masks, iou_pred
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_size": self.hidden_size,
+                "num_layers": self.num_layers,
+                "intermediate_dim": self.intermediate_dim,
+                "num_heads": self.num_heads,
+                "embedding_dim": self.embedding_dim,
+                "num_multimask_outputs": self.num_multimask_outputs,
+                "iou_head_depth": self.iou_head_depth,
+                "iou_head_hidden_dim": self.iou_head_hidden_dim,
+                "activation": self.activation,
+            }
+        )
+        return config

keras_hub/src/models/sam/sam_prompt_encoder.py

@@ -0,0 +1,336 @@
+# 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 import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.sam.sam_layers import (
+    RandomFrequencyPositionalEmbeddings,
+)
+
+
+@keras_hub_export("keras_hub.layers.SAMPromptEncoder")
+class SAMPromptEncoder(keras.layers.Layer):
+    """Prompt Encoder for the Segment Anything Model (SAM).
+
+    The prompt encoder generates encodings for three types of prompts:
+    - Point prompts: Points on the image along with a label indicating whether
+        the point is in the foreground (part of the mask) or in the background
+        (not a part of the mask).
+    - Box prompts: A batch of bounding boxes with format [(x1, y1), (x2, y2)]
+        used to determine the location of the masks in the image.
+    - Masks: An input mask can be passed to refine the positional embeddings
+        for the output mask.
+
+    First, the point prompts and box prompts are concatenated and positional
+    encodings are generated using random spatial frequencies. A point is
+    represented as the sum of a positional encoding of the point's location
+    and one of two learned embeddings that indicate if the point is either in
+    the foreground or background. A box is represented by an embedding pair:
+    (1) the positional encoding of its top-left corner summed with a learned
+    embedding representing "top-left corner" and
+    (2) the same structure but using a learned embedding indicating
+    "bottom-right corner".
+    The box and point encodings are referred to as "prompt_sparse encodings"
+    If a mask prompt is passed, a convolutional neural net is used to
+    downscale it to generate "dense encodings". If no mask prompt is passed,
+    an embedding layer is used instead to generate a "no mask" embedding.
+
+
+    Args:
+        hidden_size: int, optional. The number of features in the output
+            embeddings. Defaults to `256`.
+        image_embedding_size: int, optional. The number of features in the
+            image embeddings generated by an image encoder. Defaults to
+            `(64, 64)`.
+        input_image_size: tuple[int], optional. A tuple of the height and
+            width of the image being prompted. Defaults to `(1024, 1024)`.
+        mask_in_channels: int, optional. The number of channels of the mask
+            prompt. Defaults to `16`.
+        activation: str, optional. The activation to use in the mask
+            downscaler neural net. Defaults to `"gelu"`.
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden_size=256,
+        image_embedding_size=(64, 64),
+        input_image_size=(1024, 1024),
+        mask_in_channels=16,
+        activation="gelu",
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.image_embedding_size = image_embedding_size
+        self.input_image_size = input_image_size
+        self.mask_in_channels = mask_in_channels
+        self.activation = activation
+
+        self.positional_embedding_layer = RandomFrequencyPositionalEmbeddings(
+            num_positional_features=self.hidden_size // 2, scale=1
+        )
+
+        self.foreground_point_embed = keras.layers.Embedding(
+            1, hidden_size, name="foreground_point_embed"
+        )
+        self.background_point_embed = keras.layers.Embedding(
+            1, hidden_size, name="background_point_embed"
+        )
+        self.top_left_corner_embed = keras.layers.Embedding(
+            1, hidden_size, name="top_left_corner_embed"
+        )
+        self.bottom_right_corner_embed = keras.layers.Embedding(
+            1, hidden_size, name="bottom_right_corner_embed"
+        )
+        self.not_a_point_embed = keras.layers.Embedding(
+            1, hidden_size, name="not_a_point_embed"
+        )
+
+        self.mask_downscaler = keras.models.Sequential(
+            [
+                keras.layers.Conv2D(
+                    mask_in_channels // 4, kernel_size=2, strides=2
+                ),
+                keras.layers.LayerNormalization(epsilon=1e-6),
+                keras.layers.Activation(activation),
+                keras.layers.Conv2D(mask_in_channels, kernel_size=2, strides=2),
+                keras.layers.LayerNormalization(epsilon=1e-6),
+                keras.layers.Activation(activation),
+                keras.layers.Conv2D(hidden_size, kernel_size=1),
+            ],
+            name="mask_downscaler",
+        )
+        self.no_mask_embed = keras.layers.Embedding(
+            1, hidden_size, name="no_mask_embed"
+        )
+
+    def build(
+        self,
+        points_shape=None,
+        labels_shape=None,
+        boxes_shape=None,
+        masks_shape=None,
+    ):
+        self.positional_embedding_layer.build()
+        for layer in [
+            self.foreground_point_embed,
+            self.background_point_embed,
+            self.top_left_corner_embed,
+            self.bottom_right_corner_embed,
+            self.not_a_point_embed,
+            self.no_mask_embed,
+        ]:
+            layer.build([None])
+        self.mask_downscaler.build(
+            [
+                None,
+                4 * self.image_embedding_size[0],
+                4 * self.image_embedding_size[1],
+                1,
+            ]
+        )
+        self.built = True
+
+    def compute_output_shape(
+        self,
+        points_shape=None,
+        labels_shape=None,
+        boxes_shape=None,
+        masks_shape=None,
+    ):
+        batch_size = None
+        for shape in (points_shape, labels_shape, boxes_shape, masks_shape):
+            if shape is not None:
+                batch_size = shape[0]
+                break
+        return {
+            "prompt_sparse_embeddings": (
+                batch_size,
+                None,
+                self.hidden_size,
+            ),
+            "prompt_dense_embeddings": (
+                batch_size,
+                self.image_embedding_size[0],
+                self.image_embedding_size[1],
+                self.hidden_size,
+            ),
+            "prompt_dense_positional_embeddings": (
+                batch_size,
+                self.image_embedding_size[0],
+                self.image_embedding_size[1],
+                self.hidden_size,
+            ),
+        }
+
+    def _embed_points(self, points, labels):
+        points = points + 0.5
+        indices = ops.arange(1, dtype="int32")
+
+        point_embeddings = self.positional_embedding_layer.encode_coordinates(
+            points, self.input_image_size
+        )
+        labels = ops.broadcast_to(
+            labels[..., None], ops.shape(point_embeddings)
+        )
+        point_embeddings = ops.where(
+            labels == 0,
+            point_embeddings + self.background_point_embed(indices),
+            point_embeddings + self.foreground_point_embed(indices),
+        )
+        point_embeddings = ops.where(
+            labels == -1,
+            self.not_a_point_embed(indices),
+            point_embeddings,
+        )
+        return point_embeddings
+
+    def _embed_box(self, box):
+        shape = ops.shape(box)
+        batch_size, N = shape[0], shape[1]
+        box = box + 0.5
+        indices = ops.arange(1, dtype="int32")
+        corner_embedding = self.positional_embedding_layer.encode_coordinates(
+            box, self.input_image_size
+        )
+        top_left_embedding = corner_embedding[
+            :, :, 0, :
+        ] + self.top_left_corner_embed(indices)
+        bottom_right_embedding = corner_embedding[
+            :, :, 1, :
+        ] + self.bottom_right_corner_embed(indices)
+        corner_embedding = ops.stack(
+            [top_left_embedding, bottom_right_embedding], axis=2
+        )
+        return ops.reshape(
+            corner_embedding, (batch_size, N * 2, self.hidden_size)
+        )
+
+    def _embed_mask(self, mask):
+        mask_embedding = self.mask_downscaler(mask)
+        return mask_embedding
+
+    def call(
+        self, images=None, points=None, labels=None, boxes=None, masks=None
+    ):
+        # Get the batch shape based on any arbitrary input, because batch
+        # shapes must all match.
+        valid_inputs = [
+            x for x in (points, labels, boxes, masks) if x is not None
+        ]
+
+        batch_size = ops.shape(valid_inputs[0])[0]
+        if points is None:
+            points = ops.zeros((batch_size, 0, 2))
+        if labels is None:
+            labels = ops.zeros((batch_size, 0))
+        if boxes is None:
+            boxes = ops.zeros((batch_size, 0, 2, 2))
+        if masks is None:
+            masks = ops.zeros((batch_size, 0, 256, 256, 1))
+
+        # Compute point embeddings
+        point_embeddings = self._embed_points(points, labels)
+
+        # Compute box embeddings
+        box_embeddings = self._embed_box(boxes)
+
+        # Concatenate both into a sparse embeddings tensor
+        sparse_embeddings = ops.concatenate(
+            [point_embeddings, box_embeddings], axis=1
+        )
+
+        # Compute the mask embeddings
+        def _no_mask_embed():
+            reshaped_embed = ops.reshape(
+                self.no_mask_embed(ops.arange(1, dtype="int32")),
+                (1, 1, 1, self.hidden_size),
+            )
+            broadcasted_embed = ops.broadcast_to(
+                reshaped_embed,
+                shape=(
+                    batch_size,
+                    self.image_embedding_size[0],
+                    self.image_embedding_size[1],
+                    self.hidden_size,
+                ),
+            )
+            return broadcasted_embed
+
+        def _maybe_input_mask_embed():
+            # Keras passes the masks as concrete tensors for both the
+            # true and false functions to build the output shape. So, we
+            # need to handle the case when 0 size masks is passed and
+            # dispatch the call to `_no_mask_embed`. Note that we can't call
+            # the lambda directly since the inputs are bound to different
+            # values when called with concrete values.
+            if masks.shape[1] == 0:
+                return ops.broadcast_to(
+                    ops.reshape(
+                        self.no_mask_embed(ops.arange(1, dtype="int32")),
+                        (1, 1, 1, self.hidden_size),
+                    ),
+                    shape=(
+                        batch_size,
+                        self.image_embedding_size[0],
+                        self.image_embedding_size[1],
+                        self.hidden_size,
+                    ),
+                )
+            shape = ops.shape(masks)
+            BM, N, height, width, channels = (
+                shape[0],
+                shape[1],
+                shape[2],
+                shape[3],
+                shape[4],
+            )
+            return self._embed_mask(
+                ops.reshape(masks, (BM * N, height, width, channels))
+            )
+
+        dense_embeddings = ops.cond(
+            ops.equal(ops.size(masks), 0),
+            _no_mask_embed,
+            _maybe_input_mask_embed,
+        )
+
+        # Compute the dense positional embeddings
+        prompt_dense_positional_embeddings = (
+            self.positional_embedding_layer.encode_image(
+                self.image_embedding_size
+            )[None, ...]
+        )
+
+        return {
+            "prompt_sparse_embeddings": sparse_embeddings,
+            "prompt_dense_embeddings": dense_embeddings,
+            "prompt_dense_positional_embeddings": prompt_dense_positional_embeddings,
+        }
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_size": self.hidden_size,
+                "image_embedding_size": self.image_embedding_size,
+                "input_image_size": self.input_image_size,
+                "mask_in_channels": self.mask_in_channels,
+                "activation": self.activation,
+            }
+        )
+        return config