Made label data optional for inference and adopted other required changes (#2183)

* Made label data optional for inference and adopted other required changes

* address comments

* Fix format issues

* update version numbers for newly generated checkpoints

Author: laxmareddyp

keras_hub/src/models/segformer/segformer_image_segmenter_preprocessor.py

@@ -1,5 +1,3 @@
-import keras
-
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.image_segmenter_preprocessor import (
     ImageSegmenterPreprocessor,
@@ -8,25 +6,9 @@
 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)
-            if y is not None:
-                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_image_segmenter_tests.py

@@ -1,37 +1,45 @@
 import numpy as np
 import pytest
-from keras import ops
 
 from keras_hub.src.models.mit.mit_backbone import MiTBackbone
 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_image_segmenter_preprocessor import (  # noqa: E501
+    SegFormerImageSegmenterPreprocessor,
+)
 from keras_hub.src.tests.test_case import TestCase
 
 
 class SegFormerTest(TestCase):
     def setUp(self):
         image_encoder = MiTBackbone(
-            depths=[2, 2],
-            image_shape=(224, 224, 3),
+            layerwise_depths=[2, 2],
+            image_shape=(32, 32, 3),
             hidden_dims=[32, 64],
             num_layers=2,
-            blockwise_num_heads=[1, 2],
-            blockwise_sr_ratios=[8, 4],
+            layerwise_num_heads=[1, 2],
+            layerwise_sr_ratios=[8, 4],
             max_drop_path_rate=0.1,
-            patch_sizes=[7, 3],
-            strides=[4, 2],
+            layerwise_patch_sizes=[7, 3],
+            layerwise_strides=[4, 2],
         )
         projection_filters = 256
+        self.preprocessor = SegFormerImageSegmenterPreprocessor()
         self.backbone = SegFormerBackbone(
             image_encoder=image_encoder, projection_filters=projection_filters
         )
 
-        self.input_size = 224
-        self.input_data = ops.ones((2, self.input_size, self.input_size, 3))
+        self.input_size = 32
+        self.input_data = np.ones((2, self.input_size, self.input_size, 3))
+        self.label_data = np.ones((2, self.input_size, self.input_size, 4))
 
-        self.init_kwargs = {"backbone": self.backbone, "num_classes": 4}
+        self.init_kwargs = {
+            "backbone": self.backbone,
+            "num_classes": 4,
+            "preprocessor": self.preprocessor,
+        }
 
     def test_segformer_segmenter_construction(self):
         SegFormerImageSegmenter(backbone=self.backbone, num_classes=4)
@@ -42,19 +50,19 @@ def test_segformer_call(self):
             backbone=self.backbone, num_classes=4
         )
 
-        images = np.random.uniform(size=(2, 224, 224, 4))
+        images = np.random.uniform(size=(2, 32, 32, 3))
         segformer_output = segformer(images)
         segformer_predict = segformer.predict(images)
 
-        assert segformer_output.shape == images.shape
-        assert segformer_predict.shape == images.shape
+        self.assertAllEqual(segformer_output.shape, (2, 32, 32, 4))
+        self.assertAllEqual(segformer_predict.shape, (2, 32, 32, 4))
 
     def test_task(self):
         self.run_task_test(
             cls=SegFormerImageSegmenter,
             init_kwargs={**self.init_kwargs},
-            train_data=self.input_data,
-            expected_output_shape=(2, 224, 224),
+            train_data=(self.input_data, self.label_data),
+            expected_output_shape=(2, 32, 32, 4),
         )
 
     @pytest.mark.large

keras_hub/src/models/segformer/segformer_presets.py

@@ -10,7 +10,7 @@
             "params": 3719027,
             "path": "segformer_b0",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b0_ade20k_512/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b0_ade20k_512/3",
     },
     "segformer_b1_ade20k_512": {
         "metadata": {
@@ -21,7 +21,7 @@
             "params": 13682643,
             "path": "segformer_b1",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b1_ade20k_512/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b1_ade20k_512/5",
     },
     "segformer_b2_ade20k_512": {
         "metadata": {
@@ -32,7 +32,7 @@
             "params": 24727507,
             "path": "segformer_b2",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b2_ade20k_512/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b2_ade20k_512/3",
     },
     "segformer_b3_ade20k_512": {
         "metadata": {
@@ -43,7 +43,7 @@
             "params": 44603347,
             "path": "segformer_b3",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b3_ade20k_512/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b3_ade20k_512/3",
     },
     "segformer_b4_ade20k_512": {
         "metadata": {
@@ -54,7 +54,7 @@
             "params": 61373907,
             "path": "segformer_b4",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b4_ade20k_512/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b4_ade20k_512/3",
     },
     "segformer_b5_ade20k_640": {
         "metadata": {
@@ -65,7 +65,7 @@
             "params": 81974227,
             "path": "segformer_b5",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b5_ade20k_640/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b5_ade20k_640/3",
     },
     "segformer_b0_cityscapes_1024": {
         "metadata": {
@@ -76,7 +76,7 @@
             "params": 3719027,
             "path": "segformer_b0",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b0_cityscapes_1024/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b0_cityscapes_1024/3",
     },
     "segformer_b1_cityscapes_1024": {
         "metadata": {
@@ -87,7 +87,7 @@
             "params": 13682643,
             "path": "segformer_b1",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b1_ade20k_512/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b1_ade20k_512/1",
     },
     "segformer_b2_cityscapes_1024": {
         "metadata": {
@@ -98,7 +98,7 @@
             "params": 24727507,
             "path": "segformer_b2",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b2_cityscapes_1024/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b2_cityscapes_1024/3",
     },
     "segformer_b3_cityscapes_1024": {
         "metadata": {
@@ -109,7 +109,7 @@
             "params": 44603347,
             "path": "segformer_b3",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b3_cityscapes_1024/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b3_cityscapes_1024/3",
     },
     "segformer_b4_cityscapes_1024": {
         "metadata": {
@@ -120,7 +120,7 @@
             "params": 61373907,
             "path": "segformer_b4",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b4_cityscapes_1024/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b4_cityscapes_1024/3",
     },
     "segformer_b5_cityscapes_1024": {
         "metadata": {
@@ -131,6 +131,6 @@
             "params": 81974227,
             "path": "segformer_b5",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b5_cityscapes_1024/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b5_cityscapes_1024/3",
     },
 }

tools/checkpoint_conversion/convert_segformer_checkpoints.py

@@ -1,6 +1,6 @@
 # Usage example
-# python tools/checkpoint_conversion/convert_mix_transformer.py \
-#     --preset "B0_ade_512"
+# python tools/checkpoint_conversion/convert_segformer_checkpoints.py \
+#     --preset "b0_ade20k_512"
 
 import numpy as np
 from absl import app
@@ -94,7 +94,17 @@ def main(_):
     )
     num_classes = 150 if "ade20k" in FLAGS.preset else 19
 
-    preprocessor = SegFormerImageSegmenterPreprocessor()
+    image_converter = keras_hub.layers.SegFormerImageConverter(
+        height=resolution,
+        width=resolution,
+        scale=[
+            1.0 / (0.229 * 255.0),
+            1.0 / (0.224 * 255.0),
+            1.0 / (0.225 * 255.0),
+        ],
+        offset=[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225],
+    )
+    preprocessor = SegFormerImageSegmenterPreprocessor(image_converter)
     segformer_segmenter = keras_hub.models.SegFormerImageSegmenter(
         backbone=segformer_backbone,
         num_classes=num_classes,