keras-hub 0.25.1__py3-none-any.whl → 0.26.0.dev0__py3-none-any.whl

keras_hub/src/models/sam3/sam3_detr_encoder.py

@@ -0,0 +1,293 @@
+from keras import layers
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.sam3.sam3_layers import SAM3MLP
+from keras_hub.src.models.sam3.sam3_layers import SAM3Attention
+from keras_hub.src.models.sam3.sam3_utils import create_bidirectional_mask
+
+
+class SAM3DetrEncoderLayer(layers.Layer):
+    def __init__(
+        self,
+        hidden_dim,
+        intermediate_dim,
+        num_heads,
+        hidden_activation="relu",
+        dropout_rate=0.0,
+        layer_norm_epsilon=1e-6,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dim = int(hidden_dim)
+        self.intermediate_dim = int(intermediate_dim)
+        self.num_heads = int(num_heads)
+        self.hidden_activation = hidden_activation
+        self.dropout_rate = float(dropout_rate)
+        self.layer_norm_epsilon = float(layer_norm_epsilon)
+
+        self.layer_norm1 = layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="layer_norm1",
+        )
+        self.self_attn = SAM3Attention(
+            hidden_dim=self.hidden_dim,
+            num_heads=self.num_heads,
+            dtype=self.dtype_policy,
+            name="self_attn",
+        )
+        self.dropout = layers.Dropout(
+            self.dropout_rate, dtype=self.dtype_policy, name="dropout"
+        )
+        self.cross_attn = SAM3Attention(
+            hidden_dim=self.hidden_dim,
+            num_heads=self.num_heads,
+            dtype=self.dtype_policy,
+            name="cross_attn",
+        )
+        self.layer_norm2 = layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="layer_norm2",
+        )
+        self.mlp = SAM3MLP(
+            hidden_dim=self.hidden_dim,
+            intermediate_dim=self.intermediate_dim,
+            activation=self.hidden_activation,
+            dropout_rate=self.dropout_rate,
+            dtype=self.dtype_policy,
+            name="mlp",
+        )
+        self.layer_norm3 = layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="layer_norm3",
+        )
+
+    def build(
+        self,
+        vision_feats_shape,
+        prompt_feats_shape,
+        vision_pos_encodings_shape,
+        prompt_cross_attn_masks_shape,
+    ):
+        self.layer_norm1.build(vision_feats_shape)
+        self.self_attn.build(
+            vision_feats_shape, vision_feats_shape, vision_feats_shape
+        )
+        self.dropout.build(vision_feats_shape)
+        self.layer_norm2.build(vision_feats_shape)
+        self.cross_attn.build(
+            vision_feats_shape, prompt_feats_shape, prompt_feats_shape
+        )
+        self.layer_norm3.build(vision_feats_shape)
+        self.mlp.build(vision_feats_shape)
+
+    def call(
+        self,
+        vision_feats,
+        prompt_feats,
+        vision_pos_encodings,
+        prompt_cross_attn_masks=None,
+        training=None,
+    ):
+        residual = vision_feats
+        hidden_states = self.layer_norm1(vision_feats, training=training)
+        hidden_states_with_pos = ops.add(hidden_states, vision_pos_encodings)
+        hidden_states = self.self_attn(
+            query=hidden_states_with_pos,
+            key=hidden_states_with_pos,
+            value=hidden_states,
+            training=training,
+        )
+        hidden_states = ops.add(
+            self.dropout(hidden_states, training=training), residual
+        )
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states, training=training)
+        hidden_states = self.cross_attn(
+            query=hidden_states,
+            key=prompt_feats,
+            value=prompt_feats,
+            attention_mask=prompt_cross_attn_masks,
+            training=training,
+        )
+
+        hidden_states = ops.add(
+            self.dropout(hidden_states, training=training), residual
+        )
+
+        residual = hidden_states
+        hidden_states = self.layer_norm3(hidden_states, training=training)
+        hidden_states = self.mlp(hidden_states, training=training)
+        hidden_states = ops.add(
+            self.dropout(hidden_states, training=training), residual
+        )
+        return hidden_states
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "num_heads": self.num_heads,
+                "hidden_activation": self.hidden_activation,
+                "dropout_rate": self.dropout_rate,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+            }
+        )
+        return config
+
+    def compute_output_shape(
+        self,
+        vision_feats_shape,
+        prompt_feats_shape,
+        vision_pos_encodings_shape,
+        prompt_cross_attn_masks_shape,
+    ):
+        return vision_feats_shape
+
+
+@keras_hub_export("keras_hub.layers.SAM3DetrEncoder")
+class SAM3DetrEncoder(layers.Layer):
+    """A DETR encoder for the Segment Anything Model 3 (SAM3).
+
+    This layer implements a transformer-based encoder that fuses vision and
+    prompt features. It processes flattened vision features and prompt features
+    through multiple layers of self-attention and cross-attention.
+
+    Args:
+        num_layers: int. The number of transformer layers.
+        hidden_dim: int. The hidden dimension of the transformer layers.
+        intermediate_dim: int. The dimension of the intermediate layer in the
+            transformer's MLP.
+        num_heads: int. The number of attention heads.
+        hidden_activation: str. The activation function for the transformer
+            layers. Defaults to `"relu"`.
+        dropout_rate: float. The dropout rate for the MLP and attention.
+            Defaults to `0.0`.
+        layer_norm_epsilon: float. The epsilon value for layer normalization.
+            Defaults to `1e-6`.
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        hidden_dim,
+        intermediate_dim,
+        num_heads,
+        hidden_activation="relu",
+        dropout_rate=0.0,
+        layer_norm_epsilon=1e-6,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_layers = int(num_layers)
+        self.hidden_dim = int(hidden_dim)
+        self.intermediate_dim = int(intermediate_dim)
+        self.num_heads = int(num_heads)
+        self.hidden_activation = hidden_activation
+        self.dropout_rate = float(dropout_rate)
+        self.layer_norm_epsilon = float(layer_norm_epsilon)
+
+        self.layers = [
+            SAM3DetrEncoderLayer(
+                hidden_dim=self.hidden_dim,
+                intermediate_dim=self.intermediate_dim,
+                num_heads=self.num_heads,
+                dropout_rate=self.dropout_rate,
+                hidden_activation=self.hidden_activation,
+                layer_norm_epsilon=self.layer_norm_epsilon,
+                dtype=self.dtype_policy,
+                name=f"layer_{i}",
+            )
+            for i in range(self.num_layers)
+        ]
+
+    def build(
+        self,
+        vision_features_shape,
+        text_features_shape,
+        vision_pos_embeds_shape,
+        text_masks_shape,
+    ):
+        self.height = int(vision_features_shape[1])
+        self.width = int(vision_features_shape[2])
+        feature_flattened_shape = [
+            vision_features_shape[0],
+            vision_features_shape[1] * vision_features_shape[2],
+            vision_features_shape[-1],
+        ]
+        for layer in self.layers:
+            layer.build(
+                feature_flattened_shape,
+                text_features_shape,
+                feature_flattened_shape,
+                None,
+            )
+
+    def call(
+        self,
+        vision_features,
+        text_features,
+        vision_pos_embeds,
+        text_masks,
+        training=None,
+    ):
+        # Flatten multi-level features for encoder processing.
+        batch_size = ops.shape(vision_features)[0]
+        hidden_dim = ops.shape(vision_features)[-1]
+        features_flattened = ops.reshape(
+            vision_features,
+            (batch_size, self.height * self.width, hidden_dim),
+        )
+        pos_embeds_flattened = ops.reshape(
+            vision_pos_embeds,
+            (batch_size, self.height * self.width, hidden_dim),
+        )
+
+        prompt_cross_attn_masks = create_bidirectional_mask(
+            features_flattened, text_masks
+        )
+        hidden_states = features_flattened
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                prompt_feats=text_features,
+                vision_pos_encodings=pos_embeds_flattened,
+                prompt_cross_attn_masks=prompt_cross_attn_masks,
+                training=training,
+            )
+        return hidden_states, pos_embeds_flattened
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_layers": self.num_layers,
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "num_heads": self.num_heads,
+                "hidden_activation": self.hidden_activation,
+                "dropout_rate": self.dropout_rate,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+            }
+        )
+        return config
+
+    def compute_output_shape(
+        self,
+        vision_features_shape,
+        text_features_shape,
+        vision_pos_embeds_shape,
+        text_masks_shape,
+    ):
+        features_flattened_shape = [
+            vision_features_shape[0],
+            vision_features_shape[1] * vision_features_shape[2],
+            vision_features_shape[-1],
+        ]
+        return features_flattened_shape, features_flattened_shape

keras_hub/src/models/sam3/sam3_dot_product_scoring.py

@@ -0,0 +1,120 @@
+import numpy as np
+from keras import layers
+from keras import ops
+
+from keras_hub.src.models.sam3.sam3_layers import SAM3DecoderMLP
+
+
+class SAM3DotProductScoring(layers.Layer):
+    def __init__(
+        self,
+        hidden_dim,
+        intermediate_dim,
+        dropout_rate=0.0,
+        layer_norm_epsilon=1e-6,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dim = int(hidden_dim)
+        self.intermediate_dim = int(intermediate_dim)
+        self.dropout_rate = float(dropout_rate)
+        self.layer_norm_epsilon = float(layer_norm_epsilon)
+
+        self.text_mlp = SAM3DecoderMLP(
+            num_layers=2,
+            hidden_dim=self.intermediate_dim,
+            output_dim=self.hidden_dim,
+            dtype=self.dtype_policy,
+            name="text_mlp",
+        )
+        self.text_mlp_dropout = layers.Dropout(
+            self.dropout_rate, dtype=self.dtype_policy, name="text_mlp_dropout"
+        )
+        self.text_mlp_out_norm = layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="text_mlp_out_norm",
+        )
+
+        # Projections for text and query features.
+        self.text_proj = layers.Dense(
+            self.hidden_dim, dtype=self.dtype_policy, name="text_proj"
+        )
+        self.query_proj = layers.Dense(
+            self.hidden_dim, dtype=self.dtype_policy, name="query_proj"
+        )
+
+        # Scale factor for dot product.
+        self.scale = float(1.0 / np.sqrt(self.hidden_dim))
+
+        # Clamping to avoid numerical issues.
+        self.clamp_max_val = 12.0
+
+    def build(
+        self, decoder_hidden_states_shape, text_features_shape, text_masks_shape
+    ):
+        self.text_mlp.build(text_features_shape)
+        self.text_mlp_dropout.build(text_features_shape)
+        self.text_mlp_out_norm.build(text_features_shape)
+        pooled_text_shape = [text_features_shape[0], text_features_shape[-1]]
+        self.text_proj.build(pooled_text_shape)
+        self.query_proj.build(decoder_hidden_states_shape)
+
+    def _pool_text_features(self, text_features, text_mask=None):
+        if text_mask is None:
+            # No padding, simple mean.
+            return ops.mean(text_features, axis=1)
+
+        is_valid = ops.expand_dims(
+            ops.cast(text_mask, text_features.dtype), axis=-1
+        )
+        # Count valid tokens per batch.
+        num_valid = ops.maximum(ops.sum(is_valid, axis=1), 1.0)
+        # Mean pool only over valid tokens.
+        return ops.divide(
+            ops.sum(ops.multiply(text_features, is_valid), axis=1), num_valid
+        )
+
+    def call(
+        self,
+        decoder_hidden_states,
+        text_features,
+        text_masks=None,
+        training=None,
+    ):
+        orig_text_features = text_features
+        text_features = self.text_mlp(text_features, training=training)
+        text_features = self.text_mlp_dropout(text_features, training=training)
+        text_features = ops.add(text_features, orig_text_features)
+        text_features = self.text_mlp_out_norm(text_features, training=training)
+
+        pooled_text = self._pool_text_features(text_features, text_masks)
+
+        proj_text = self.text_proj(pooled_text, training=training)
+        proj_queries = self.query_proj(decoder_hidden_states, training=training)
+
+        proj_text = ops.expand_dims(proj_text, axis=-1)
+        scores = ops.matmul(proj_queries, ops.expand_dims(proj_text, axis=1))
+        scores = ops.multiply(scores, self.scale)
+        scores = ops.clip(scores, -self.clamp_max_val, self.clamp_max_val)
+        return scores
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "dropout_rate": self.dropout_rate,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+            }
+        )
+        return config
+
+    def compute_output_shape(
+        self, decoder_hidden_states_shape, text_features_shape, text_masks_shape
+    ):
+        batch_size = decoder_hidden_states_shape[0]
+        num_layers = decoder_hidden_states_shape[1]
+        num_queries = decoder_hidden_states_shape[2]
+        return [batch_size, num_layers, num_queries, 1]