Add support for decoding uint16 images in tfexample_decoder.py.

PiperOrigin-RevId: 667708306 Change-Id: I0218d51a6035959ec41019de22d49e36e820d68d
google-research · Aug 26, 2024 · 3aafefa · 3aafefa
1 parent 8d78490
commit 3aafefa
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/tf_slim/data/tfexample_decoder.py b/tf_slim/data/tfexample_decoder.py
@@ -423,8 +423,11 @@ def _decode(self, image_buffer, image_format):
     def decode_image():
       """Decodes a image based on the headers."""
       return math_ops.cast(
-          image_ops.decode_image(image_buffer, channels=self._channels),
-          self._dtype)
+          image_ops.decode_image(
+              image_buffer, channels=self._channels, dtype=self._dtype
+          ),
+          self._dtype,
+      )
 
     def decode_jpeg():
       """Decodes a jpeg image with specified '_dct_method'."""