Add SAM3 Promptable Concept Segmentation (PCS) model (keras-team#2534)

* Add SAM3 PCS model.

* Rename to PromptableConcept.

* Update docstrings and add post-processing.

* Remove testing code. Rename to SAM3PromptableConceptImageSegmenterPreprocessor.

* Fix quantization.

* Address comments.

* Add custom image shape support.

Author: james77777778

keras_hub/api/layers/__init__.py

@@ -138,6 +138,27 @@
 from keras_hub.src.models.sam.sam_prompt_encoder import (
     SAMPromptEncoder as SAMPromptEncoder,
 )
+from keras_hub.src.models.sam3.sam3_detr_decoder import (
+    SAM3DetrDecoder as SAM3DetrDecoder,
+)
+from keras_hub.src.models.sam3.sam3_detr_encoder import (
+    SAM3DetrEncoder as SAM3DetrEncoder,
+)
+from keras_hub.src.models.sam3.sam3_geometry_encoder import (
+    SAM3GeometryEncoder as SAM3GeometryEncoder,
+)
+from keras_hub.src.models.sam3.sam3_image_converter import (
+    SAM3ImageConverter as SAM3ImageConverter,
+)
+from keras_hub.src.models.sam3.sam3_mask_decoder import (
+    SAM3MaskDecoder as SAM3MaskDecoder,
+)
+from keras_hub.src.models.sam3.sam3_text_encoder import (
+    SAM3TextEncoder as SAM3TextEncoder,
+)
+from keras_hub.src.models.sam3.sam3_vision_encoder import (
+    SAM3VisionEncoder as SAM3VisionEncoder,
+)
 from keras_hub.src.models.segformer.segformer_image_converter import (
     SegFormerImageConverter as SegFormerImageConverter,
 )

keras_hub/api/models/__init__.py

@@ -651,6 +651,18 @@
 from keras_hub.src.models.sam.sam_image_segmenter_preprocessor import (
     SAMImageSegmenterPreprocessor as SAMImageSegmenterPreprocessor,
 )
+from keras_hub.src.models.sam3.sam3_pc_backbone import (
+    SAM3PromptableConceptBackbone as SAM3PromptableConceptBackbone,
+)
+from keras_hub.src.models.sam3.sam3_pc_image_segmenter import (
+    SAM3PromptableConceptImageSegmenter as SAM3PromptableConceptImageSegmenter,
+)
+from keras_hub.src.models.sam3.sam3_pc_image_segmenter_preprocessor import (
+    SAM3PromptableConceptImageSegmenterPreprocessor as SAM3PromptableConceptImageSegmenterPreprocessor,
+)
+from keras_hub.src.models.sam3.sam3_tokenizer import (
+    SAM3Tokenizer as SAM3Tokenizer,
+)
 from keras_hub.src.models.segformer.segformer_backbone import (
     SegFormerBackbone as SegFormerBackbone,
 )

keras_hub/api/tokenizers/__init__.py

@@ -96,6 +96,9 @@
 from keras_hub.src.models.rwkv7.rwkv7_tokenizer import (
     RWKVTokenizer as RWKVTokenizer,
 )
+from keras_hub.src.models.sam3.sam3_tokenizer import (
+    SAM3Tokenizer as SAM3Tokenizer,
+)
 from keras_hub.src.models.siglip.siglip_tokenizer import (
     SigLIPTokenizer as SigLIPTokenizer,
 )

keras_hub/src/models/sam/sam_backbone.py

@@ -16,7 +16,11 @@ class SAMBackbone(Backbone):
         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.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for the models computations and weights. Note that some
+            computations, such as softmax and layer normalization will always
+            be done in float32 precision regardless of dtype. Defaults to
+            `bfloat16`.
 
     Example:
     ```python

keras_hub/src/models/sam3/__init__.py

@@ -0,0 +1,7 @@
+from keras_hub.src.models.sam3.sam3_pc_backbone import (
+    SAM3PromptableConceptBackbone,
+)
+from keras_hub.src.models.sam3.sam3_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, SAM3PromptableConceptBackbone)

keras_hub/src/models/sam3/roi_align.py

@@ -0,0 +1,222 @@
+from keras import backend
+from keras import ops
+
+
+def _bilinear_interpolate(
+    feature_maps, roi_batch_ind, y, x, ymask, xmask, height, width, hidden_dim
+):
+    feature_maps_dtype = backend.standardize_dtype(feature_maps.dtype)
+    y = ops.maximum(y, 0.0)
+    x = ops.maximum(x, 0.0)
+    y_low = ops.cast(y, "int32")
+    x_low = ops.cast(x, "int32")
+    y_high = ops.where(
+        ops.greater_equal(y_low, height - 1), height - 1, y_low + 1
+    )
+    y_low = ops.where(ops.greater_equal(y_low, height - 1), height - 1, y_low)
+    y = ops.where(
+        ops.greater_equal(y_low, height - 1),
+        ops.cast(y, dtype=feature_maps_dtype),
+        y,
+    )
+
+    x_high = ops.where(
+        ops.greater_equal(x_low, width - 1), width - 1, x_low + 1
+    )
+    x_low = ops.where(ops.greater_equal(x_low, width - 1), width - 1, x_low)
+    x = ops.where(
+        ops.greater_equal(x_low, width - 1),
+        ops.cast(x, dtype=feature_maps_dtype),
+        x,
+    )
+
+    ly = ops.subtract(y, y_low)
+    lx = ops.subtract(x, x_low)
+    hy = ops.subtract(1.0, ly)
+    hx = ops.subtract(1.0, lx)
+
+    def masked_index(y, x):
+        y = ops.where(ymask[:, None, :], y, 0)
+        x = ops.where(xmask[:, None, :], x, 0)
+        batch_idx = roi_batch_ind[:, None, None, None, None, None]
+        channel_idx = ops.arange(hidden_dim)[None, None, None, None, None, :]
+        y_idx = y[:, :, None, :, None, None]
+        x_idx = x[:, None, :, None, :, None]
+
+        if backend.backend() == "tensorflow":
+            import tensorflow as tf
+
+            # Explicitly broadcast indices to the same shape for XLA
+            # compatibility
+            common_zero = ops.zeros_like(
+                batch_idx + y_idx + x_idx + channel_idx
+            )
+            batch_idx = batch_idx + common_zero
+            y_idx = y_idx + common_zero
+            x_idx = ops.transpose(
+                ops.transpose(x_idx, (0, 2, 1, 4, 3, 5)) + common_zero,
+                (0, 2, 1, 4, 3, 5),
+            )
+            channel_idx = channel_idx + common_zero
+            indices = ops.stack([batch_idx, y_idx, x_idx, channel_idx], axis=-1)
+            indices = ops.cast(indices, "int32")
+            return tf.gather_nd(feature_maps, indices)
+        else:
+            return feature_maps[
+                batch_idx,
+                y_idx,
+                x_idx,
+                channel_idx,
+            ]
+
+    v1 = masked_index(y_low, x_low)
+    v2 = masked_index(y_low, x_high)
+    v3 = masked_index(y_high, x_low)
+    v4 = masked_index(y_high, x_high)
+
+    def outer_prod(y, x):
+        return ops.multiply(
+            y[:, :, None, :, None, None], x[:, None, :, None, :, None]
+        )
+
+    w1 = outer_prod(hy, hx)
+    w2 = outer_prod(hy, lx)
+    w3 = outer_prod(ly, hx)
+    w4 = outer_prod(ly, lx)
+
+    val = ops.add(
+        ops.add(ops.multiply(w1, v1), ops.multiply(w2, v2)),
+        ops.add(ops.multiply(w3, v3), ops.multiply(w4, v4)),
+    )
+    return val
+
+
+def roi_align_torch(
+    feature_maps,
+    rois,
+    output_size,
+    spatial_scale=1.0,
+    aligned=False,
+):
+    import torchvision
+
+    dtype = backend.standardize_dtype(feature_maps.dtype)
+    need_cast = False
+    if dtype == "bfloat16":
+        # torchvision.ops.roi_align does not support bfloat16.
+        feature_maps = ops.cast(feature_maps, "float32")
+        rois = ops.cast(rois, "float32")
+        need_cast = True
+
+    output = ops.transpose(
+        torchvision.ops.roi_align(
+            ops.transpose(feature_maps, (0, 3, 1, 2)),
+            rois,
+            output_size,
+            spatial_scale=spatial_scale,
+            aligned=aligned,
+        ),
+        (0, 2, 3, 1),
+    )
+    if need_cast:
+        output = ops.cast(output, dtype)
+    return output
+
+
+def roi_align(
+    feature_maps,
+    rois,
+    output_size,
+    height,
+    width,
+    hidden_dim,
+    spatial_scale=1.0,
+    aligned=False,
+):
+    # Use torchvision's optimized roi_align implementation.
+    if backend.backend() == "torch":
+        return roi_align_torch(
+            feature_maps,
+            rois,
+            output_size,
+            spatial_scale=spatial_scale,
+            aligned=aligned,
+        )
+
+    original_dtype = backend.standardize_dtype(feature_maps.dtype)
+    out_h, out_w = output_size[0], output_size[1]
+
+    feature_maps = ops.cast(feature_maps, "float32")
+    rois = ops.cast(rois, "float32")
+
+    ph = ops.arange(out_h, dtype="float32")
+    pw = ops.arange(out_w, dtype="float32")
+
+    # input: [N, C, H, W]
+    # rois: [K, 5]
+
+    roi_batch_ind = ops.cast(rois[:, 0], "int32")
+    offset = 0.5 if aligned else 0.0
+    roi_start_w = ops.subtract(ops.multiply(rois[:, 1], spatial_scale), offset)
+    roi_start_h = ops.subtract(ops.multiply(rois[:, 2], spatial_scale), offset)
+    roi_end_w = ops.subtract(ops.multiply(rois[:, 3], spatial_scale), offset)
+    roi_end_h = ops.subtract(ops.multiply(rois[:, 4], spatial_scale), offset)
+
+    roi_width = ops.subtract(roi_end_w, roi_start_w)
+    roi_height = ops.subtract(roi_end_h, roi_start_h)
+    if not aligned:
+        roi_width = ops.maximum(roi_width, 1.0)
+        roi_height = ops.maximum(roi_height, 1.0)
+
+    bin_size_h = ops.divide(roi_height, out_h)
+    bin_size_w = ops.divide(roi_width, out_w)
+
+    roi_bin_grid_h = ops.ceil(ops.divide(roi_height, out_h))
+    roi_bin_grid_w = ops.ceil(ops.divide(roi_width, out_w))
+
+    count = ops.maximum(ops.multiply(roi_bin_grid_h, roi_bin_grid_w), 1.0)
+    iy = ops.arange(height, dtype="float32")
+    ix = ops.arange(width, dtype="float32")
+    ymask = ops.less(iy[None, :], roi_bin_grid_h[:, None])
+    xmask = ops.less(ix[None, :], roi_bin_grid_w[:, None])
+
+    def from_k(t):
+        return t[:, None, None]
+
+    y = ops.add(
+        ops.add(
+            from_k(roi_start_h),
+            ops.multiply(ph[None, :, None], from_k(bin_size_h)),
+        ),
+        ops.multiply(
+            ops.cast(ops.add(iy[None, None, :], 0.5), dtype="float32"),
+            from_k(ops.divide(bin_size_h, roi_bin_grid_h)),
+        ),
+    )
+    x = ops.add(
+        ops.add(
+            from_k(roi_start_w),
+            ops.multiply(pw[None, :, None], from_k(bin_size_w)),
+        ),
+        ops.multiply(
+            ops.cast(ops.add(ix[None, None, :], 0.5), dtype="float32"),
+            from_k(ops.divide(bin_size_w, roi_bin_grid_w)),
+        ),
+    )
+    val = _bilinear_interpolate(
+        feature_maps,
+        roi_batch_ind,
+        y,
+        x,
+        ymask,
+        xmask,
+        height,
+        width,
+        hidden_dim,
+    )
+    val = ops.where(ymask[:, None, None, :, None, None], val, 0.0)
+    val = ops.where(xmask[:, None, None, None, :, None], val, 0.0)
+
+    output = ops.sum(val, axis=(3, 4))
+    output = ops.divide(output, count[:, None, None, None])
+    return ops.cast(output, original_dtype)