keras-hub-nightly 0.16.1.dev202410200345__py3-none-any.whl → 0.19.0.dev202412070351__py3-none-any.whl

keras_hub/src/models/sam/sam_image_segmenter.py

@@ -51,9 +51,9 @@ class SAMImageSegmenter(ImageSegmenter):
             (batch_size, 0, image_size, image_size, 1)
         ),
     }
-    # todo: update preset name
-    sam = keras_hub.models.SAMImageSegmenter.from_preset(`sam_base`)
-    sam(input_data)
+    sam = keras_hub.models.SAMImageSegmenter.from_preset('sam_base_sa1b')
+    outputs = sam.predict(input_data)
+    masks, iou_pred = outputs["masks"], outputs["iou_pred"]
     ```
 
     Load segment anything image segmenter with custom backbone
@@ -65,7 +65,7 @@ class SAMImageSegmenter(ImageSegmenter):
         (batch_size, image_size, image_size, 3),
         dtype="float32",
     )
-    image_encoder = ViTDetBackbone(
+    image_encoder = keras_hub.models.ViTDetBackbone(
         hidden_size=16,
         num_layers=16,
         intermediate_dim=16 * 4,
@@ -76,7 +76,7 @@ class SAMImageSegmenter(ImageSegmenter):
         window_size=2,
         image_shape=(image_size, image_size, 3),
     )
-    prompt_encoder = SAMPromptEncoder(
+    prompt_encoder = keras_hub.layers.SAMPromptEncoder(
         hidden_size=8,
         image_embedding_size=(8, 8),
         input_image_size=(
@@ -85,7 +85,7 @@ class SAMImageSegmenter(ImageSegmenter):
         ),
         mask_in_channels=16,
     )
-    mask_decoder = SAMMaskDecoder(
+    mask_decoder = keras_hub.layers.SAMMaskDecoder(
         num_layers=2,
         hidden_size=8,
         intermediate_dim=32,
@@ -95,13 +95,12 @@ class SAMImageSegmenter(ImageSegmenter):
         iou_head_depth=3,
         iou_head_hidden_dim=8,
     )
-    backbone = SAMBackbone(
+    backbone = keras_hub.models.SAMBackbone(
         image_encoder=image_encoder,
         prompt_encoder=prompt_encoder,
         mask_decoder=mask_decoder,
-        image_shape=(image_size, image_size, 3),
     )
-    sam = SAMImageSegmenter(
+    sam = keras_hub.models.SAMImageSegmenter(
         backbone=backbone
     )
     ```
@@ -115,7 +114,7 @@ class SAMImageSegmenter(ImageSegmenter):
     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:
+    # Prepare an input dictionary:
     inputs = {
         "images": image,
         "points": points,

keras_hub/src/models/sam/sam_presets.py

@@ -5,9 +5,7 @@ backbone_presets = {
         "metadata": {
             "description": ("The base SAM model trained on the SA1B dataset."),
             "params": 93735728,
-            "official_name": "SAMImageSegmenter",
             "path": "sam",
-            "model_card": "https://arxiv.org/abs/2304.02643",
         },
         "kaggle_handle": "kaggle://keras/sam/keras/sam_base_sa1b/4",
     },
@@ -15,9 +13,7 @@ backbone_presets = {
         "metadata": {
             "description": ("The large SAM model trained on the SA1B dataset."),
             "params": 641090864,
-            "official_name": "SAMImageSegmenter",
             "path": "sam",
-            "model_card": "https://arxiv.org/abs/2304.02643",
         },
         "kaggle_handle": "kaggle://keras/sam/keras/sam_large_sa1b/4",
     },
@@ -25,9 +21,7 @@ backbone_presets = {
         "metadata": {
             "description": ("The huge SAM model trained on the SA1B dataset."),
             "params": 312343088,
-            "official_name": "SAMImageSegmenter",
             "path": "sam",
-            "model_card": "https://arxiv.org/abs/2304.02643",
         },
         "kaggle_handle": "kaggle://keras/sam/keras/sam_huge_sa1b/4",
     },

keras_hub/src/models/segformer/__init__.py

@@ -0,0 +1,8 @@
+from keras_hub.src.models.segformer.segformer_backbone import SegFormerBackbone
+from keras_hub.src.models.segformer.segformer_image_segmenter import (
+    SegFormerImageSegmenter,
+)
+from keras_hub.src.models.segformer.segformer_presets import presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(presets, SegFormerImageSegmenter)

keras_hub/src/models/segformer/segformer_backbone.py

@@ -0,0 +1,163 @@
+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.SegFormerBackbone")
+class SegFormerBackbone(Backbone):
+    """A Keras model implementing the SegFormer architecture for semantic segmentation.
+
+    This class implements the majority of the SegFormer architecture described in
+    [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers]
+    (https://arxiv.org/abs/2105.15203) and [based on the TensorFlow implementation from DeepVision]
+    (https://github.com/DavidLandup0/deepvision/tree/main/deepvision/models/segmentation/segformer).
+
+    SegFormers are meant to be used with the MixTransformer (MiT) encoder family, and
+    and use a very lightweight all-MLP decoder head.
+
+    The MiT encoder uses a hierarchical transformer which outputs features at multiple scales,
+    similar to that of the hierarchical outputs typically associated with CNNs.
+
+    Args:
+        image_encoder: `keras.Model`. The backbone network for the model that is
+            used as a feature extractor for the SegFormer encoder.
+            Should be used with the MiT backbone model
+            (`keras_hub.models.MiTBackbone`) which was created
+            specifically for SegFormers.
+        num_classes: int, the number of classes for the detection model,
+            including the background class.
+        projection_filters: int, number of filters in the
+            convolution layer projecting the concatenated features into
+            a segmentation map. Defaults to 256`.
+
+    Example:
+
+    Using the class with a custom `backbone`:
+
+    ```python
+    import keras_hub
+
+    backbone = keras_hub.models.MiTBackbone(
+        depths=[2, 2, 2, 2],
+        image_shape=(224, 224, 3),
+        hidden_dims=[32, 64, 160, 256],
+        num_layers=4,
+        blockwise_num_heads=[1, 2, 5, 8],
+        blockwise_sr_ratios=[8, 4, 2, 1],
+        max_drop_path_rate=0.1,
+        patch_sizes=[7, 3, 3, 3],
+        strides=[4, 2, 2, 2],
+    )
+
+    segformer_backbone = keras_hub.models.SegFormerBackbone(image_encoder=backbone, projection_filters=256)
+    ```
+
+    Using the class with a preset `backbone`:
+
+    ```python
+    import keras_hub
+
+    backbone = keras_hub.models.MiTBackbone.from_preset("mit_b0_ade20k_512")
+    segformer_backbone = keras_hub.models.SegFormerBackbone(image_encoder=backbone, projection_filters=256)
+    ```
+
+    """
+
+    def __init__(
+        self,
+        image_encoder,
+        projection_filters,
+        **kwargs,
+    ):
+        if not isinstance(image_encoder, keras.layers.Layer) or not isinstance(
+            image_encoder, keras.Model
+        ):
+            raise ValueError(
+                "Argument `image_encoder` must be a `keras.layers.Layer` instance "
+                f" or `keras.Model`. Received instead "
+                f"image_encoder={image_encoder} (of type {type(image_encoder)})."
+            )
+
+        # === Layers ===
+        inputs = keras.layers.Input(shape=image_encoder.input.shape[1:])
+
+        self.feature_extractor = keras.Model(
+            image_encoder.inputs, image_encoder.pyramid_outputs
+        )
+
+        features = self.feature_extractor(inputs)
+        # Get height and width of level one output
+        _, height, width, _ = features["P1"].shape
+
+        self.mlp_blocks = []
+
+        for feature_dim, feature in zip(image_encoder.hidden_dims, features):
+            self.mlp_blocks.append(
+                keras.layers.Dense(
+                    projection_filters, name=f"linear_{feature_dim}"
+                )
+            )
+
+        self.resizing = keras.layers.Resizing(
+            height, width, interpolation="bilinear"
+        )
+        self.concat = keras.layers.Concatenate(axis=-1)
+        self.linear_fuse = keras.Sequential(
+            [
+                keras.layers.Conv2D(
+                    filters=projection_filters, kernel_size=1, use_bias=False
+                ),
+                keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.9),
+                keras.layers.Activation("relu"),
+            ]
+        )
+
+        # === Functional Model ===
+        # Project all multi-level outputs onto
+        # the same dimensionality and feature map shape
+        multi_layer_outs = []
+        for index, (feature_dim, feature) in enumerate(
+            zip(image_encoder.hidden_dims, features)
+        ):
+            out = self.mlp_blocks[index](features[feature])
+            out = self.resizing(out)
+            multi_layer_outs.append(out)
+
+        # Concat now-equal feature maps
+        concatenated_outs = self.concat(multi_layer_outs[::-1])
+
+        # Fuse concatenated features into a segmentation map
+        seg = self.linear_fuse(concatenated_outs)
+
+        super().__init__(
+            inputs=inputs,
+            outputs=seg,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.projection_filters = projection_filters
+        self.image_encoder = image_encoder
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "projection_filters": self.projection_filters,
+                "image_encoder": keras.saving.serialize_keras_object(
+                    self.image_encoder
+                ),
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        if "image_encoder" in config and isinstance(
+            config["image_encoder"], dict
+        ):
+            config["image_encoder"] = keras.layers.deserialize(
+                config["image_encoder"]
+            )
+        return super().from_config(config)

keras_hub/src/models/segformer/segformer_image_converter.py

@@ -0,0 +1,8 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
+from keras_hub.src.models.segformer.segformer_backbone import SegFormerBackbone
+
+
+@keras_hub_export("keras_hub.layers.SegFormerImageConverter")
+class SegFormerImageConverter(ImageConverter):
+    backbone_cls = SegFormerBackbone

keras_hub/src/models/segformer/segformer_image_segmenter.py

@@ -0,0 +1,171 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.image_segmenter import ImageSegmenter
+from keras_hub.src.models.segformer.segformer_backbone import SegFormerBackbone
+from keras_hub.src.models.segformer.segformer_image_segmenter_preprocessor import (
+    SegFormerImageSegmenterPreprocessor,
+)
+
+
+@keras_hub_export("keras_hub.models.SegFormerImageSegmenter")
+class SegFormerImageSegmenter(ImageSegmenter):
+    """A Keras model implementing the SegFormer architecture for semantic segmentation.
+
+    This class implements the segmentation head of the SegFormer architecture described in
+    [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers]
+    (https://arxiv.org/abs/2105.15203) and [based on the TensorFlow implementation from DeepVision]
+    (https://github.com/DavidLandup0/deepvision/tree/main/deepvision/models/segmentation/segformer).
+
+    SegFormers are meant to be used with the MixTransformer (MiT) encoder family, and
+    and use a very lightweight all-MLP decoder head.
+
+    The MiT encoder uses a hierarchical transformer which outputs features at multiple scales,
+    similar to that of the hierarchical outputs typically associated with CNNs.
+
+    Args:
+        image_encoder: `keras.Model`. The backbone network for the model that is
+            used as a feature extractor for the SegFormer encoder.
+            It is *intended* to be used only with the MiT backbone model
+            (`keras_hub.models.MiTBackbone`) which was created
+            specifically for SegFormers.
+            Alternatively, can be a `keras_hub.models.Backbone` a model subclassing
+            `keras_hub.models.FeaturePyramidBackbone`, or a `keras.Model`
+            that has a `pyramid_outputs` property which is
+            a dictionary with keys "P2", "P3", "P4", and "P5" and layer names as values.
+        num_classes: int, the number of classes for the detection model,
+            including the background class.
+        projection_filters: int, number of filters in the
+            convolution layer projecting the concatenated features into
+            a segmentation map. Defaults to 256`.
+
+
+    Example:
+
+    Using presets:
+
+    ```python
+    import keras_hub
+    import numpy as np
+
+    segmenter = keras_hub.models.SegFormerImageSegmenter.from_preset("segformer_b0_ade20k_512")
+
+    images = np.random.rand(1, 512, 512, 3)
+    segformer(images)
+    ```
+
+    Using the SegFormer backbone:
+
+    ```python
+    encoder = keras_hub.models.MiTBackbone.from_preset("mit_b0_ade20k_512")
+    backbone = keras_hub.models.SegFormerBackbone(image_encoder=encoder, projection_filters=256)
+    ```
+
+    Using the SegFormer backbone with a custom encoder:
+
+    ```python
+    import keras
+    import keras_hub
+    import numpy as np
+
+    images = np.ones(shape=(1, 96, 96, 3))
+    labels = np.zeros(shape=(1, 96, 96, 1))
+
+    encoder = keras_hub.models.MiTBackbone(
+        depths=[2, 2, 2, 2],
+        image_shape=(96, 96, 3),
+        hidden_dims=[32, 64, 160, 256],
+        num_layers=4,
+        blockwise_num_heads=[1, 2, 5, 8],
+        blockwise_sr_ratios=[8, 4, 2, 1],
+        max_drop_path_rate=0.1,
+        patch_sizes=[7, 3, 3, 3],
+        strides=[4, 2, 2, 2],
+    )
+
+    backbone = keras_hub.models.SegFormerBackbone(image_encoder=encoder, projection_filters=256)
+    segformer = keras_hub.models.SegFormerImageSegmenter(backbone=backbone, num_classes=4)
+
+    segformer(images)
+    ```
+
+    Using the segmentor class with a preset backbone:
+
+    ```python
+    import keras_hub
+
+    image_encoder = keras_hub.models.MiTBackbone.from_preset("mit_b0_ade20k_512")
+    backbone = keras_hub.models.SegFormerBackbone(image_encoder=encoder, projection_filters=256)
+    segformer = keras_hub.models.SegFormerImageSegmenter(backbone=backbone, num_classes=4)
+    ```
+    """
+
+    backbone_cls = SegFormerBackbone
+    preprocessor_cls = SegFormerImageSegmenterPreprocessor
+
+    def __init__(
+        self,
+        backbone,
+        num_classes,
+        preprocessor=None,
+        **kwargs,
+    ):
+        if not isinstance(backbone, keras.layers.Layer) or not isinstance(
+            backbone, keras.Model
+        ):
+            raise ValueError(
+                "Argument `backbone` must be a `keras.layers.Layer` instance "
+                f" or `keras.Model`. Received instead "
+                f"backbone={backbone} (of type {type(backbone)})."
+            )
+
+        # === Layers ===
+        inputs = backbone.input
+
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.dropout = keras.layers.Dropout(0.1)
+        self.output_segmentation_head = keras.layers.Conv2D(
+            filters=num_classes, kernel_size=1, strides=1
+        )
+        self.resizing = keras.layers.Resizing(
+            height=inputs.shape[1],
+            width=inputs.shape[2],
+            interpolation="bilinear",
+        )
+
+        # === Functional Model ===
+        x = self.backbone(inputs)
+        x = self.dropout(x)
+        x = self.output_segmentation_head(x)
+        output = self.resizing(x)
+
+        super().__init__(
+            inputs=inputs,
+            outputs=output,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.num_classes = num_classes
+        self.backbone = backbone
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_classes": self.num_classes,
+                "backbone": keras.saving.serialize_keras_object(self.backbone),
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        if "image_encoder" in config and isinstance(
+            config["image_encoder"], dict
+        ):
+            config["image_encoder"] = keras.layers.deserialize(
+                config["image_encoder"]
+            )
+        return super().from_config(config)

keras_hub/src/models/segformer/segformer_image_segmenter_preprocessor.py

@@ -0,0 +1,31 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.image_segmenter_preprocessor import (
+    ImageSegmenterPreprocessor,
+)
+from keras_hub.src.models.segformer.segformer_backbone import SegFormerBackbone
+from keras_hub.src.models.segformer.segformer_image_converter import (
+    SegFormerImageConverter,
+)
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+IMAGENET_DEFAULT_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_DEFAULT_STD = [0.229, 0.224, 0.225]
+
+
+@keras_hub_export("keras_hub.models.SegFormerImageSegmenterPreprocessor")
+class SegFormerImageSegmenterPreprocessor(ImageSegmenterPreprocessor):
+    backbone_cls = SegFormerBackbone
+    image_converter_cls = SegFormerImageConverter
+
+    @preprocessing_function
+    def call(self, x, y=None, sample_weight=None):
+        if self.image_converter:
+            x = self.image_converter(x)
+            y = self.image_converter(y)
+
+        x = x / 255
+        x = (x - IMAGENET_DEFAULT_MEAN) / IMAGENET_DEFAULT_STD
+
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)

keras_hub/src/models/segformer/segformer_presets.py

@@ -0,0 +1,124 @@
+"""SegFormer model preset configurations."""
+
+presets = {
+    "segformer_b0_ade20k_512": {
+        "metadata": {
+            "description": (
+                "SegFormer model with MiTB0 backbone fine-tuned on ADE20k in 512x512 resolution."
+            ),
+            "params": 3719027,
+            "path": "segformer_b0",
+        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b0_ade20k_512/2",
+    },
+    "segformer_b1_ade20k_512": {
+        "metadata": {
+            "description": (
+                "SegFormer model with MiTB1 backbone fine-tuned on ADE20k in 512x512 resolution."
+            ),
+            "params": 13682643,
+            "path": "segformer_b1",
+        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b1_ade20k_512/2",
+    },
+    "segformer_b2_ade20k_512": {
+        "metadata": {
+            "description": (
+                "SegFormer model with MiTB2 backbone fine-tuned on ADE20k in 512x512 resolution."
+            ),
+            "params": 24727507,
+            "path": "segformer_b2",
+        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b2_ade20k_512/2",
+    },
+    "segformer_b3_ade20k_512": {
+        "metadata": {
+            "description": (
+                "SegFormer model with MiTB3 backbone fine-tuned on ADE20k in 512x512 resolution."
+            ),
+            "params": 44603347,
+            "path": "segformer_b3",
+        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b3_ade20k_512/2",
+    },
+    "segformer_b4_ade20k_512": {
+        "metadata": {
+            "description": (
+                "SegFormer model with MiTB4 backbone fine-tuned on ADE20k in 512x512 resolution."
+            ),
+            "params": 61373907,
+            "path": "segformer_b4",
+        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b4_ade20k_512/2",
+    },
+    "segformer_b5_ade20k_640": {
+        "metadata": {
+            "description": (
+                "SegFormer model with MiTB5 backbone fine-tuned on ADE20k in 640x640 resolution."
+            ),
+            "params": 81974227,
+            "path": "segformer_b5",
+        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b5_ade20k_640/2",
+    },
+    "segformer_b0_cityscapes_1024": {
+        "metadata": {
+            "description": (
+                "SegFormer model with MiTB0 backbone fine-tuned on Cityscapes in 1024x1024 resolution."
+            ),
+            "params": 3719027,
+            "path": "segformer_b0",
+        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b0_cityscapes_1024/2",
+    },
+    "segformer_b1_cityscapes_1024": {
+        "metadata": {
+            "description": (
+                "SegFormer model with MiTB1 backbone fine-tuned on Cityscapes in 1024x1024 resolution."
+            ),
+            "params": 13682643,
+            "path": "segformer_b1",
+        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b1_ade20k_512/2",
+    },
+    "segformer_b2_cityscapes_1024": {
+        "metadata": {
+            "description": (
+                "SegFormer model with MiTB2 backbone fine-tuned on Cityscapes in 1024x1024 resolution."
+            ),
+            "params": 24727507,
+            "path": "segformer_b2",
+        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b2_cityscapes_1024/2",
+    },
+    "segformer_b3_cityscapes_1024": {
+        "metadata": {
+            "description": (
+                "SegFormer model with MiTB3 backbone fine-tuned on Cityscapes in 1024x1024 resolution."
+            ),
+            "params": 44603347,
+            "path": "segformer_b3",
+        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b3_cityscapes_1024/2",
+    },
+    "segformer_b4_cityscapes_1024": {
+        "metadata": {
+            "description": (
+                "SegFormer model with MiTB4 backbone fine-tuned on Cityscapes in 1024x1024 resolution."
+            ),
+            "params": 61373907,
+            "path": "segformer_b4",
+        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b4_cityscapes_1024/2",
+    },
+    "segformer_b5_cityscapes_1024": {
+        "metadata": {
+            "description": (
+                "SegFormer model with MiTB5 backbone fine-tuned on Cityscapes in 1024x1024 resolution."
+            ),
+            "params": 81974227,
+            "path": "segformer_b5",
+        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b5_cityscapes_1024/2",
+    },
+}