Fix typing annotation

PiperOrigin-RevId: 424086207

Author: Conchylicultor

tensorflow_datasets/core/utils/image_utils.py

@@ -45,7 +45,7 @@ def decode_image(image_bytes: bytes) -> np.ndarray:
   return runner.run(tf.image.decode_image, image_bytes)
 
 
-def png_to_jpeg(image_bytes: bytes, quality: int = 100) -> np.ndarray:
+def png_to_jpeg(image_bytes: bytes, quality: int = 100) -> bytes:
   """Converts PNG image (bytes or str) to JPEG (bytes)."""
   runner = _get_runner()
   decode_fn = lambda img: tf.image.decode_png(img, channels=3)
@@ -54,7 +54,7 @@ def png_to_jpeg(image_bytes: bytes, quality: int = 100) -> np.ndarray:
   return runner.run(fn, image)
 
 
-def jpeg_cmyk_to_rgb(image_bytes: bytes, quality: int = 100) -> np.ndarray:
+def jpeg_cmyk_to_rgb(image_bytes: bytes, quality: int = 100) -> bytes:
   """Converts JPEG CMYK image (bytes) to RGB JPEG (bytes)."""
   runner = _get_runner()
   image = runner.run(tf.image.decode_jpeg, image_bytes)

tensorflow_datasets/core/utils/image_utils_test.py

@@ -52,12 +52,14 @@ def test_png_to_jpeg(self):
   def test_png_4chan_to_jpeg(self):
     image = self._get_image('6pixels_4chan.png')
     jpeg = image_utils.png_to_jpeg(image)
+    self.assertIsInstance(jpeg, bytes)
     image_np_jpeg = image_utils.decode_image(jpeg)
     np.testing.assert_array_equal(image_np_jpeg, SIX_PIXELS_JPEG)
 
   def test_jpeg_cmyk_to_rgb(self):
     image = self._get_image('6pixels_cmyk.jpeg')
     new_image = image_utils.jpeg_cmyk_to_rgb(image, quality=100)
+    self.assertIsInstance(new_image, bytes)
     self.assertNotEqual(image, new_image)
     # Converting between color systems is not bijective, so high rtol.
     original_np_image = image_utils.decode_image(image)

tensorflow_datasets/image_classification/imagenet.py

@@ -259,10 +259,9 @@ def _fix_image(self, image_fname, image):
     if self.version < '3.0.0':
       return image
     if image_fname in CMYK_IMAGES:
-      image = io.BytesIO(
-          tfds.core.utils.jpeg_cmyk_to_rgb(image.read()).tobytes())
+      image = io.BytesIO(tfds.core.utils.jpeg_cmyk_to_rgb(image.read()))
     elif image_fname in PNG_IMAGES:
-      image = io.BytesIO(tfds.core.utils.png_to_jpeg(image.read()).tobytes())
+      image = io.BytesIO(tfds.core.utils.png_to_jpeg(image.read()))
     return image
 
   def _generate_examples(self,