Fix dtype bug in image converter (#2147)

Author: abheesht17

keras_hub/src/layers/preprocessing/image_converter.py

@@ -280,6 +280,8 @@ def call(self, inputs):
         return inputs
 
     def _expand_non_channel_dims(self, value, inputs):
+        input_dtype = keras.backend.standardize_dtype(inputs.dtype)
+
         unbatched = len(ops.shape(inputs)) == 3
         channels_first = self.data_format == "channels_first"
         if unbatched:
@@ -294,9 +296,10 @@ def _expand_non_channel_dims(self, value, inputs):
             # device (potentially GPU) after preprocessing.
             if keras.backend.backend() == "torch" and self.image_size is None:
                 return ops.expand_dims(value, broadcast_dims).cpu()
-            return ops.expand_dims(value, broadcast_dims)
+            expanded = ops.expand_dims(value, broadcast_dims)
+            return ops.cast(expanded, input_dtype)
         else:
-            return np.expand_dims(value, broadcast_dims)
+            return np.expand_dims(value, broadcast_dims).astype(input_dtype)
 
     def get_config(self):
         config = super().get_config()

keras_hub/src/layers/preprocessing/image_converter_test.py

@@ -35,6 +35,21 @@ def test_unbatched(self):
         self.assertAllClose(outputs[:, :, 1], np.ones((4, 4)) * 0.301569)
         self.assertAllClose(outputs[:, :, 2], np.ones((4, 4)) * 0.852353)
 
+    def test_bfloat16_input(self):
+        converter = ImageConverter(
+            image_size=(4, 4),
+            scale=(1.0 / 255.0, 0.8 / 255.0, 1.2 / 255.0),
+            offset=(0.2, -0.1, 0.25),
+            dtype="bfloat16",
+        )
+        inputs = ops.ones((10, 10, 3)) * 128
+        inputs = ops.cast(inputs, "bfloat16")
+        outputs = converter(inputs)
+        self.assertEqual(ops.shape(outputs), (4, 4, 3))
+        self.assertAllClose(outputs[:, :, 0], np.ones((4, 4)) * 0.703125)
+        self.assertAllClose(outputs[:, :, 1], np.ones((4, 4)) * 0.302734)
+        self.assertAllClose(outputs[:, :, 2], np.ones((4, 4)) * 0.851562)
+
     @parameterized.parameters(
         (True, False),
         (False, True),