[Mix Transformer] Add Presets for MiTB0...MiTB5

keras_hub/api/layers/__init__.py

@@ -40,6 +40,9 @@
 from keras_hub.src.models.densenet.densenet_image_converter import (
     DenseNetImageConverter,
 )
+from keras_hub.src.models.mix_transformer.mix_transformer_image_converter import (
+    MiTImageConverter,
+)
 from keras_hub.src.models.pali_gemma.pali_gemma_image_converter import (
     PaliGemmaImageConverter,
 )

keras_hub/api/models/__init__.py

@@ -206,6 +206,9 @@
 from keras_hub.src.models.mix_transformer.mix_transformer_classifier import (
     MiTImageClassifier,
 )
+from keras_hub.src.models.mix_transformer.mix_transformer_classifier_preprocessor import (
+    MiTImageClassifierPreprocessor,
+)
 from keras_hub.src.models.mobilenet.mobilenet_backbone import MobileNetBackbone
 from keras_hub.src.models.mobilenet.mobilenet_image_classifier import (
     MobileNetImageClassifier,

keras_hub/src/models/mix_transformer/__init__.py

@@ -0,0 +1,12 @@
+from keras_hub.src.models.mix_transformer.mix_transformer_backbone import (
+    MiTBackbone,
+)
+from keras_hub.src.models.mix_transformer.mix_transformer_classifier import (
+    MiTImageClassifier,
+)
+from keras_hub.src.models.mix_transformer.mix_transformer_presets import (
+    backbone_presets,
+)
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, MiTBackbone)

keras_hub/src/models/mix_transformer/mix_transformer_classifier.py

@@ -3,8 +3,12 @@
 from keras_hub.src.models.mix_transformer.mix_transformer_backbone import (
     MiTBackbone,
 )
+from keras_hub.src.models.mix_transformer.mix_transformer_classifier_preprocessor import (
+    MiTImageClassifierPreprocessor,
+)
 
 
 @keras_hub_export("keras_hub.models.MiTImageClassifier")
 class MiTImageClassifier(ImageClassifier):
     backbone_cls = MiTBackbone
+    preprocessor_cls = MiTImageClassifierPreprocessor

keras_hub/src/models/mix_transformer/mix_transformer_classifier_preprocessor.py

@@ -0,0 +1,16 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.image_classifier_preprocessor import (
+    ImageClassifierPreprocessor,
+)
+from keras_hub.src.models.mix_transformer.mix_transformer_backbone import (
+    MiTBackbone,
+)
+from keras_hub.src.models.mix_transformer.mix_transformer_image_converter import (
+    MiTImageConverter,
+)
+
+
+@keras_hub_export("keras_hub.models.MiTImageClassifierPreprocessor")
+class MiTImageClassifierPreprocessor(ImageClassifierPreprocessor):
+    backbone_cls = MiTBackbone
+    image_converter_cls = MiTImageConverter

keras_hub/src/models/mix_transformer/mix_transformer_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.mix_transformer import MiTBackbone
+
+
+@keras_hub_export("keras_hub.layers.MiTImageConverter")
+class MiTImageConverter(ImageConverter):
+    backbone_cls = MiTBackbone

keras_hub/src/models/mix_transformer/mix_transformer_layers.py

@@ -28,19 +28,23 @@ def __init__(self, project_dim=32, patch_size=7, stride=4, **kwargs):
         self.patch_size = patch_size
         self.stride = stride
 
+        padding_size = self.patch_size // 2
+
+        self.padding = keras.layers.ZeroPadding2D(
+            padding=(padding_size, padding_size)
+        )
         self.proj = keras.layers.Conv2D(
             filters=project_dim,
             kernel_size=patch_size,
             strides=stride,
-            padding="same",
+            padding="valid",
         )
-        self.norm = keras.layers.LayerNormalization()
+        self.norm = keras.layers.LayerNormalization(epsilon=1e-5)
 
     def call(self, x):
+        x = self.padding(x)
         x = self.proj(x)
-        # B, H, W, C
-        shape = x.shape
-        x = ops.reshape(x, (-1, shape[1] * shape[2], shape[3]))
+        x = ops.reshape(x, (-1, x.shape[1] * x.shape[2], x.shape[3]))
         x = self.norm(x)
         return x
 

keras_hub/src/models/mix_transformer/mix_transformer_presets.py

@@ -0,0 +1,151 @@
+# 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.
+"""MiT model preset configurations."""
+
+backbone_presets_with_weights = {
+    "mit_b0_ade20k_512": {
+        "metadata": {
+            "description": (
+                "MiT (MixTransformer) model with 8 transformer blocks."
+            ),
+            "params": 3321962,
+            "official_name": "MiT",
+            "path": "mit",
+        },
+        "kaggle_handle": "kaggle://kerashub/mix-transformer/keras/mit_b0_ade20k_512",
+    },
+    "mit_b1_ade20k_512": {
+        "metadata": {
+            "description": (
+                "MiT (MixTransformer) model with 8 transformer blocks."
+            ),
+            "params": 13156554,
+            "official_name": "MiT",
+            "path": "mit",
+        },
+        "kaggle_handle": "kaggle://kerashub/mix-transformer/keras/mit_b1_ade20k_512",
+    },
+    "mit_b2_ade20k_512": {
+        "metadata": {
+            "description": (
+                "MiT (MixTransformer) model with 16 transformer blocks."
+            ),
+            "params": 24201418,
+            "official_name": "MiT",
+            "path": "mit",
+        },
+        "kaggle_handle": "kaggle://kerashub/mix-transformer/keras/mit_b2_ade20k_512",
+    },
+    "mit_b3_ade20k_512": {
+        "metadata": {
+            "description": (
+                "MiT (MixTransformer) model with 28 transformer blocks."
+            ),
+            "params": 44077258,
+            "official_name": "MiT",
+            "path": "mit",
+        },
+        "kaggle_handle": "kaggle://kerashub/mix-transformer/keras/mit_b3_ade20k_512",
+    },
+    "mit_b4_ade20k_512": {
+        "metadata": {
+            "description": (
+                "MiT (MixTransformer) model with 41 transformer blocks."
+            ),
+            "params": 60847818,
+            "official_name": "MiT",
+            "path": "mit",
+        },
+        "kaggle_handle": "kaggle://kerashub/mix-transformer/keras/mit_b4_ade20k_512",
+    },
+    "mit_b5_ade20k_640": {
+        "metadata": {
+            "description": (
+                "MiT (MixTransformer) model with 52 transformer blocks."
+            ),
+            "params": 81448138,
+            "official_name": "MiT",
+            "path": "mit",
+        },
+        "kaggle_handle": "kaggle://kerashub/mix-transformer/keras/mit_b5_ade20k_512",
+    },
+    "mit_b0_cityscapes_1024": {
+        "metadata": {
+            "description": (
+                "MiT (MixTransformer) model with 8 transformer blocks."
+            ),
+            "params": 3321962,
+            "official_name": "MiT",
+            "path": "mit",
+        },
+        "kaggle_handle": "kaggle://kerashub/mix-transformer/keras/mit_b0_cityscapes_1024",
+    },
+    "mit_b1_cityscapes_1024": {
+        "metadata": {
+            "description": (
+                "MiT (MixTransformer) model with 8 transformer blocks."
+            ),
+            "params": 13156554,
+            "official_name": "MiT",
+            "path": "mit",
+        },
+        "kaggle_handle": "kaggle://kerashub/mix-transformer/keras/mit_b1_cityscapes_1024",
+    },
+    "mit_b2_cityscapes_1024": {
+        "metadata": {
+            "description": (
+                "MiT (MixTransformer) model with 16 transformer blocks."
+            ),
+            "params": 24201418,
+            "official_name": "MiT",
+            "path": "mit",
+        },
+        "kaggle_handle": "kaggle://kerashub/mix-transformer/keras/mit_b2_cityscapes_1024",
+    },
+    "mit_b3_cityscapes_1024": {
+        "metadata": {
+            "description": (
+                "MiT (MixTransformer) model with 28 transformer blocks."
+            ),
+            "params": 44077258,
+            "official_name": "MiT",
+            "path": "mit",
+        },
+        "kaggle_handle": "kaggle://kerashub/mix-transformer/keras/mit_b3_cityscapes_1024",
+    },
+    "mit_b4_cityscapes_1024": {
+        "metadata": {
+            "description": (
+                "MiT (MixTransformer) model with 41 transformer blocks."
+            ),
+            "params": 60847818,
+            "official_name": "MiT",
+            "path": "mit",
+        },
+        "kaggle_handle": "kaggle://kerashub/mix-transformer/keras/mit_b4_cityscapes_1024",
+    },
+    "mit_b5_cityscapes_1024": {
+        "metadata": {
+            "description": (
+                "MiT (MixTransformer) model with 52 transformer blocks."
+            ),
+            "params": 81448138,
+            "official_name": "MiT",
+            "path": "mit",
+        },
+        "kaggle_handle": "kaggle://kerashub/mix-transformer/keras/mit_b5_cityscapes_1024",
+    },
+}
+
+backbone_presets = {
+    **backbone_presets_with_weights,
+}