Fix PIL image conversion for uint16 multi-channel arrays.

PiperOrigin-RevId: 911926026

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/utils/image_utils.py

@@ -177,16 +177,25 @@ def create_thumbnail(
   if use_colormap:  # Apply the colormap first as it modify the shape/dtype
     ex = apply_colormap(ex)
 
-  _, _, c = ex.shape
+  c = ex.shape[-1] if ex.ndim == 3 else 1
   postprocess = _postprocess_noop
   if c == 1:
     ex = ex.squeeze(axis=-1)
-    mode = 'L'
+    if ex.dtype == np.uint16:
+      mode = 'I;16'
+      postprocess = _postprocess_convert_rgb
+    else:
+      mode = 'L'
   elif ex.dtype == np.uint16:
-    mode = 'I;16'
+    # PIL.Image.fromarray doesn't support uint16 for >1 channels.
+    # https://github.com/python-pillow/Pillow/blob/11.0.0/src/PIL/Image.py#L3225
+    # Scale to 8-bit for visualization.
+    ex = (ex / 257).astype(np.uint8)
+    mode = None
     postprocess = _postprocess_convert_rgb
   else:
     mode = None
+    postprocess = _postprocess_convert_rgb
   img = PIL_Image.fromarray(ex, mode=mode)
   img = postprocess(img)
   if default_dimensions:

tensorflow_datasets/core/utils/image_utils_test.py

@@ -18,6 +18,7 @@
 import os
 
 import numpy as np
+import pytest
 import tensorflow as tf
 from tensorflow_datasets import testing
 from tensorflow_datasets.core.utils import image_utils
@@ -94,5 +95,26 @@ def test_apply_colormap(self):
     )
 
 
+@pytest.mark.parametrize(
+    'dtype,fill_value,expected_mode,channels,expected_shape',
+    [
+        (np.uint8, 255, 'L', 1, (32, 32)),
+        (np.uint8, 255, 'RGB', 3, (32, 32, 3)),
+        (np.uint16, 65535, 'RGB', 1, (32, 32, 3)),
+        (np.uint16, 65535, 'RGB', 3, (32, 32, 3)),
+    ],
+)
+def test_create_thumbnail_conversion(
+    dtype, fill_value, expected_mode, channels, expected_shape
+):
+  img = np.full((32, 32, channels), fill_value=fill_value, dtype=dtype)
+  thumbnail = image_utils.create_thumbnail(
+      img, use_colormap=False, default_dimensions=True
+  )
+  assert thumbnail.mode == expected_mode
+  expected_array = np.full(expected_shape, fill_value=255, dtype=np.uint8)
+  np.testing.assert_array_equal(np.asarray(thumbnail), expected_array)
+
+
 if __name__ == '__main__':
   testing.test_main()