Repository: tensorflow/datasets (python)

Issue:
Allow uint16 for png encoded image feature
Currently, `tfds.features.Image` with `png` encoding defaults to `tf.uint8` data type and cannot be customized. 
https://github.com/tensorflow/datasets/blob/96fe8c3156b43eb776456522ac739a0a6cdb86d2/tensorflow_datasets/core/features/image_feature.py#L149
However, `tf.image.decode_png` also allows `tf.uint16`, which is for instance important for medical data. One option would be to make the `dtype` for the decoding function configurable.
The encoding constraint at
https://github.com/tensorflow/datasets/blob/96fe8c3156b43eb776456522ac739a0a6cdb86d2/tensorflow_datasets/core/features/image_feature.py#L130 
should also be modified for `png` encoding to allow `np.uint8` as well as `np.uint16`.

Pull request description:
Add support for uint16 PNG images
Currently, `tfds.features.Image` with `png` encoding defaults to `tf.uint8` data type and cannot be customised.

This PR adds support for `tf.uint16` PNG encoding, which is useful for medical data or depth estimation tasks.

Fixes #213 
Closes #712

//cc @Conchylicultor

Patch:
diff --git a/tensorflow_datasets/core/features/image_feature.py b/tensorflow_datasets/core/features/image_feature.py
index 362984300..3319fecd3 100644
--- a/tensorflow_datasets/core/features/image_feature.py
+++ b/tensorflow_datasets/core/features/image_feature.py
@@ -40,6 +40,11 @@ ACCEPTABLE_CHANNELS = {
     'jpeg': (0, 1, 3),
 }
 
+ACCEPTABLE_DTYPES = {
+    'png': [tf.uint8, tf.uint16],
+    'jpeg': [tf.uint8],
+}
+
 
 class Image(feature.FeatureConnector):
   """`FeatureConnector` for images.
@@ -69,7 +74,7 @@ class Image(feature.FeatureConnector):
   """
 
   @api_utils.disallow_positional_args
-  def __init__(self, shape=None, encoding_format=None):
+  def __init__(self, shape=None, encoding_format=None, dtype=None):
     """Construct the connector.
 
     Args:
@@ -83,6 +88,8 @@ class Image(feature.FeatureConnector):
         images on disk.
         If image is loaded from {bmg,gif,jpeg,png} file, this parameter is
         ignored, and file original encoding is used.
+      dtype: tf.uint16 or tf.uint8 (default).
+        tf.uint16 can be used only with png encoding_format
 
     Raises:
       ValueError: If the shape is invalid
@@ -90,10 +97,21 @@ class Image(feature.FeatureConnector):
     self._encoding_format = None
     self._shape = None
     self._runner = None
+    self._dtype = None
 
     # Set and validate values
     self.set_encoding_format(encoding_format or 'png')
     self.set_shape(shape or (None, None, 3))
+    self.set_dtype(dtype or tf.uint8)
+
+  def set_dtype(self, dtype):
+    """Update the dtype."""
+    dtype = tf.as_dtype(dtype)
+    acceptable_dtypes = ACCEPTABLE_DTYPES[self._encoding_format]
+    if dtype not in acceptable_dtypes:
+      raise ValueError('Acceptable `dtype` for %s: %s (was %s)' % (
+          self._encoding_format, acceptable_dtypes, dtype))
+    self._dtype = dtype
 
   def set_encoding_format(self, encoding_format):
     """Update the encoding format."""
@@ -113,7 +131,7 @@ class Image(feature.FeatureConnector):
 
   def get_tensor_info(self):
     # Image is returned as a 3-d uint8 tf.Tensor.
-    return feature.TensorInfo(shape=self._shape, dtype=tf.uint8)
+    return feature.TensorInfo(shape=self._shape, dtype=self._dtype)
 
   def get_serialized_info(self):
     # Only store raw image (includes size).
@@ -123,8 +141,9 @@ class Image(feature.FeatureConnector):
     """Returns np_image encoded as jpeg or png."""
     if not self._runner:
       self._runner = utils.TFGraphRunner()
-    if np_image.dtype != np.uint8:
-      raise ValueError('Image should be uint8. Detected: %s.' % np_image.dtype)
+    if np_image.dtype != self._dtype.as_numpy_dtype:
+      raise ValueError('Image dtype should be %s. Detected: %s.' % (
+          self._dtype.as_numpy_dtype, np_image.dtype))
     utils.assert_shape_match(np_image.shape, self._shape)
     return self._runner.run(ENCODE_FN[self._encoding_format], np_image)
 
@@ -147,7 +166,7 @@ class Image(feature.FeatureConnector):
   def decode_example(self, example):
     """Reconstruct the image from the tf example."""
     img = tf.image.decode_image(
-        example, channels=self._shape[-1], dtype=tf.uint8)
+        example, channels=self._shape[-1], dtype=self._dtype)
     img.set_shape(self._shape)
     return img
 
@@ -158,6 +177,7 @@ class Image(feature.FeatureConnector):
       json.dump({
           'shape': [-1 if d is None else d for d in self._shape],
           'encoding_format': self._encoding_format,
+          'dtype': self._dtype.name,
       }, f, sort_keys=True)
 
   def load_metadata(self, data_dir, feature_name=None):
@@ -169,6 +189,7 @@ class Image(feature.FeatureConnector):
         info_data = json.load(f)
       self.set_encoding_format(info_data['encoding_format'])
       self.set_shape([None if d == -1 else d for d in info_data['shape']])
+      self.set_dtype(info_data.get('dtype', tf.uint8))
 
 
 def _get_metadata_filepath(data_dir, feature_name):


Test patch:
diff --git a/tensorflow_datasets/core/features/image_feature_test.py b/tensorflow_datasets/core/features/image_feature_test.py
index 54bb0c7db..f2d6c5c1f 100644
--- a/tensorflow_datasets/core/features/image_feature_test.py
+++ b/tensorflow_datasets/core/features/image_feature_test.py
@@ -21,6 +21,7 @@ from __future__ import print_function
 
 import os
 
+from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 from tensorflow_datasets import testing
@@ -32,23 +33,27 @@ tf.compat.v1.enable_eager_execution()
 randint = np.random.randint
 
 
-class ImageFeatureTest(testing.FeatureExpectationsTestCase):
+class ImageFeatureTest(
+    testing.FeatureExpectationsTestCase, parameterized.TestCase):
 
-  def test_images(self):
-
-    img = randint(256, size=(128, 100, 3), dtype=np.uint8)
-    img_other_shape = randint(256, size=(64, 200, 3), dtype=np.uint8)
+  @parameterized.parameters(tf.uint8, tf.uint16)
+  def test_images(self, dtype):
+    np_dtype = dtype.as_numpy_dtype
+    img = randint(256, size=(128, 100, 3), dtype=np_dtype)
+    img_other_shape = randint(256, size=(64, 200, 3), dtype=np_dtype)
     img_file_path = os.path.join(os.path.dirname(__file__),
                                  '../../testing/test_data/6pixels.png')
-    img_file_expected_content = [  # see tests_data/README.md
+    img_file_expected_content = np.array([  # see tests_data/README.md
         [[0, 255, 0], [255, 0, 0], [255, 0, 255]],
         [[0, 0, 255], [255, 255, 0], [126, 127, 128]],
-    ]
+    ])
+    if dtype == tf.uint16:
+      img_file_expected_content *= 257
 
     self.assertFeature(
-        feature=features_lib.Image(),
+        feature=features_lib.Image(dtype=dtype),
         shape=(None, None, 3),
-        dtype=tf.uint8,
+        dtype=dtype,
         tests=[
             # Numpy array
             testing.FeatureExpectationItem(
@@ -58,7 +63,7 @@ class ImageFeatureTest(testing.FeatureExpectationsTestCase):
             # File path
             testing.FeatureExpectationItem(
                 value=img_file_path,
-                expected=img_file_expected_content,
+                expected=img_file_expected_content.tolist(),
             ),
             # 'img' shape can be dynamic
             testing.FeatureExpectationItem(
@@ -69,17 +74,17 @@ class ImageFeatureTest(testing.FeatureExpectationsTestCase):
             testing.FeatureExpectationItem(
                 value=randint(256, size=(128, 128, 3), dtype=np.uint32),
                 raise_cls=ValueError,
-                raise_msg='should be uint8',
+                raise_msg='dtype should be',
             ),
             # Invalid number of dimensions
             testing.FeatureExpectationItem(
-                value=randint(256, size=(128, 128), dtype=np.uint8),
+                value=randint(256, size=(128, 128), dtype=np_dtype),
                 raise_cls=ValueError,
                 raise_msg='must have the same rank',
             ),
             # Invalid number of channels
             testing.FeatureExpectationItem(
-                value=randint(256, size=(128, 128, 1), dtype=np.uint8),
+                value=randint(256, size=(128, 128, 1), dtype=np_dtype),
                 raise_cls=ValueError,
                 raise_msg='are incompatible',
             ),