Rename mix_tranformer to mit

keras_hub/api/layers/__init__.py

@@ -40,9 +40,7 @@
 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.mit.mit_image_converter import MiTImageConverter
 from keras_hub.src.models.pali_gemma.pali_gemma_image_converter import (
     PaliGemmaImageConverter,
 )

keras_hub/api/models/__init__.py

@@ -202,13 +202,9 @@
     MistralCausalLMPreprocessor,
 )
 from keras_hub.src.models.mistral.mistral_tokenizer import MistralTokenizer
-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_classifier_preprocessor import (
+from keras_hub.src.models.mit.mit_backbone import MiTBackbone
+from keras_hub.src.models.mit.mit_classifier import MiTImageClassifier
+from keras_hub.src.models.mit.mit_classifier_preprocessor import (
     MiTImageClassifierPreprocessor,
 )
 from keras_hub.src.models.mobilenet.mobilenet_backbone import MobileNetBackbone

keras_hub/src/models/mit/__init__.py

@@ -0,0 +1,6 @@
+from keras_hub.src.models.mit.mit_backbone import MiTBackbone
+from keras_hub.src.models.mit.mit_classifier import MiTImageClassifier
+from keras_hub.src.models.mit.mit_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_backbone.py → keras_hub/src/models/mit/mit_backbone.py

@@ -4,12 +4,8 @@
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.feature_pyramid_backbone import FeaturePyramidBackbone
-from keras_hub.src.models.mix_transformer.mix_transformer_layers import (
-    HierarchicalTransformerEncoder,
-)
-from keras_hub.src.models.mix_transformer.mix_transformer_layers import (
-    OverlappingPatchingAndEmbedding,
-)
+from keras_hub.src.models.mit.mit_layers import HierarchicalTransformerEncoder
+from keras_hub.src.models.mit.mit_layers import OverlappingPatchingAndEmbedding
 
 
 @keras_hub_export("keras_hub.models.MiTBackbone")

keras_hub/src/models/mit/mit_backbone_test.py

@@ -0,0 +1,45 @@
+import numpy as np
+import pytest
+
+from keras_hub.src.models.mit.mit_backbone import MiTBackbone
+from keras_hub.src.tests.test_case import TestCase
+
+
+class MiTBackboneTest(TestCase):
+    def setUp(self):
+        self.init_kwargs = {
+            "depths": [2, 2],
+            "image_shape": (16, 16, 3),
+            "hidden_dims": [4, 8],
+            "num_layers": 2,
+            "blockwise_num_heads": [1, 2],
+            "blockwise_sr_ratios": [8, 4],
+            "max_drop_path_rate": 0.1,
+            "patch_sizes": [7, 3],
+            "strides": [4, 2],
+        }
+        self.input_size = 16
+        self.input_data = np.ones(
+            (2, self.input_size, self.input_size, 3), dtype="float32"
+        )
+
+    def test_backbone_basics(self):
+        self.run_vision_backbone_test(
+            cls=MiTBackbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            expected_output_shape=(2, 2, 2, 8),
+            expected_pyramid_output_keys=["P1", "P2"],
+            expected_pyramid_image_sizes=[(4, 4), (2, 2)],
+            run_quantization_check=False,
+            run_mixed_precision_check=False,
+            run_data_format_check=False,
+        )
+
+    @pytest.mark.large
+    def test_saved_model(self):
+        self.run_model_saving_test(
+            cls=MiTBackbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )

keras_hub/src/models/mix_transformer/mix_transformer_classifier.py → keras_hub/src/models/mit/mit_classifier.py

@@ -1,9 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.image_classifier import ImageClassifier
-from keras_hub.src.models.mix_transformer.mix_transformer_backbone import (
-    MiTBackbone,
-)
-from keras_hub.src.models.mix_transformer.mix_transformer_classifier_preprocessor import (
+from keras_hub.src.models.mit.mit_backbone import MiTBackbone
+from keras_hub.src.models.mit.mit_classifier_preprocessor import (
     MiTImageClassifierPreprocessor,
 )
 

keras_hub/src/models/mix_transformer/mix_transformer_classifier_preprocessor.py → keras_hub/src/models/mit/mit_classifier_preprocessor.py

@@ -2,12 +2,8 @@
 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,
-)
+from keras_hub.src.models.mit.mit_backbone import MiTBackbone
+from keras_hub.src.models.mit.mit_image_converter import MiTImageConverter
 
 
 @keras_hub_export("keras_hub.models.MiTImageClassifierPreprocessor")

keras_hub/src/models/mix_transformer/mix_transformer_classifier_test.py → keras_hub/src/models/mit/mit_classifier_test.py

@@ -1,12 +1,8 @@
 import numpy as np
 import pytest
 
-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.mit.mit_backbone import MiTBackbone
+from keras_hub.src.models.mit.mit_classifier import MiTImageClassifier
 from keras_hub.src.tests.test_case import TestCase
 
 

keras_hub/src/models/mix_transformer/mix_transformer_image_converter.py → keras_hub/src/models/mit/mit_image_converter.py

@@ -1,6 +1,6 @@
 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
+from keras_hub.src.models.mit import MiTBackbone
 
 
 @keras_hub_export("keras_hub.layers.MiTImageConverter")

keras_hub/src/models/mix_transformer/mix_transformer_backbone_test.py → keras_hub/src/models/mit/mix_transformer_backbone_test.py

@@ -1,9 +1,7 @@
 import numpy as np
 import pytest
 
-from keras_hub.src.models.mix_transformer.mix_transformer_backbone import (
-    MiTBackbone,
-)
+from keras_hub.src.models.mit.mit_backbone import MiTBackbone
 from keras_hub.src.tests.test_case import TestCase