Fix wrong behavior (regression) with tf.image.per_image_standardization. Partially reverting PR #28210 (906e0e3bc0). Also added test cases validating against additional data types.

There were several reasons the changes introduced in #28210 did not work:
- Converting standardized image (float) to input `image` data type using `tf.image.convert_image_dtype(..., saturate=True)` loses data (clips half of data) for unsigned data types (reported in #44983 (https://github.com/tensorflow/tensorflow/issues/44983)).
- Standardized images with zero mean will have negative values (unless all values are `0`s). `convert_image_dtype` by the API definition (50216cbe96/tensorflow/python/ops/image_ops_impl.py (L2231-L2234)), however, requires that floating input dtype be in the range [0, 1).

PiperOrigin-RevId: 348379825
Change-Id: I4af7a97b391470a2018738bc6cc573e106e7f4af

tensorflow/python/ops/image_ops_impl.py

@@ -1837,7 +1837,7 @@ def per_image_standardization(image):
       dimensions of each image.
 
   Returns:
-    A `Tensor` with the same shape and dtype as `image`.
+    A `Tensor` with the same shape as `image`.
 
   Raises:
     ValueError: if the shape of 'image' is incompatible with this function.
@@ -1846,22 +1846,18 @@ def per_image_standardization(image):
     image = ops.convert_to_tensor(image, name='image')
     image = _AssertAtLeast3DImage(image)
 
-    # Remember original dtype to so we can convert back if needed
-    orig_dtype = image.dtype
-    if orig_dtype not in [dtypes.float16, dtypes.float32]:
-      image = convert_image_dtype(image, dtypes.float32)
-
+    image = math_ops.cast(image, dtype=dtypes.float32)
     num_pixels = math_ops.reduce_prod(array_ops.shape(image)[-3:])
     image_mean = math_ops.reduce_mean(image, axis=[-1, -2, -3], keepdims=True)
 
     # Apply a minimum normalization that protects us against uniform images.
     stddev = math_ops.reduce_std(image, axis=[-1, -2, -3], keepdims=True)
-    min_stddev = math_ops.rsqrt(math_ops.cast(num_pixels, image.dtype))
+    min_stddev = math_ops.rsqrt(math_ops.cast(num_pixels, dtypes.float32))
     adjusted_stddev = math_ops.maximum(stddev, min_stddev)
 
     image -= image_mean
     image = math_ops.divide(image, adjusted_stddev, name=scope)
-    return convert_image_dtype(image, orig_dtype, saturate=True)
+    return image
 
 
 @tf_export('image.random_brightness')

tensorflow/python/ops/image_ops_test.py

@@ -1643,7 +1643,8 @@ class AdjustBrightnessTest(test_util.TensorFlowTestCase):
     self._testBrightness(x_np, y_np, delta=-10. / 255.)
 
 
-class PerImageWhiteningTest(test_util.TensorFlowTestCase):
+class PerImageWhiteningTest(test_util.TensorFlowTestCase,
+                            parameterized.TestCase):
 
   def _NumpyPerImageWhitening(self, x):
     num_pixels = np.prod(x.shape)
@@ -1656,13 +1657,19 @@ class PerImageWhiteningTest(test_util.TensorFlowTestCase):
     y /= stddev
     return y
 
-  def testBasic(self):
+  @parameterized.named_parameters([("_int8", np.int8), ("_int16", np.int16),
+                                   ("_int32", np.int32), ("_int64", np.int64),
+                                   ("_uint8", np.uint8), ("_uint16", np.uint16),
+                                   ("_uint32", np.uint32),
+                                   ("_uint64", np.uint64),
+                                   ("_float32", np.float32)])
+  def testBasic(self, data_type):
     x_shape = [13, 9, 3]
-    x_np = np.arange(0, np.prod(x_shape), dtype=np.float32).reshape(x_shape)
+    x_np = np.arange(0, np.prod(x_shape), dtype=data_type).reshape(x_shape)
     y_np = self._NumpyPerImageWhitening(x_np)
 
     with self.cached_session(use_gpu=True):
-      x = constant_op.constant(x_np, shape=x_shape)
+      x = constant_op.constant(x_np, dtype=data_type, shape=x_shape)
       y = image_ops.per_image_standardization(x)
       y_tf = self.evaluate(y)
       self.assertAllClose(y_tf, y_np, atol=1e-4)
@@ -1685,17 +1692,6 @@ class PerImageWhiteningTest(test_util.TensorFlowTestCase):
       for w_tf, w_np in zip(whiten_tf, whiten_np):
         self.assertAllClose(w_tf, w_np, atol=1e-4)
 
-  def testPreservesDtype(self):
-    imgs_npu8 = np.random.uniform(0., 255., [2, 5, 5, 3]).astype(np.uint8)
-    imgs_tfu8 = constant_op.constant(imgs_npu8)
-    whiten_tfu8 = image_ops.per_image_standardization(imgs_tfu8)
-    self.assertEqual(whiten_tfu8.dtype, dtypes.uint8)
-
-    imgs_npf16 = np.random.uniform(0., 255., [2, 5, 5, 3]).astype(np.float16)
-    imgs_tff16 = constant_op.constant(imgs_npf16)
-    whiten_tff16 = image_ops.per_image_standardization(imgs_tff16)
-    self.assertEqual(whiten_tff16.dtype, dtypes.float16)
-
 
 class CropToBoundingBoxTest(test_util.TensorFlowTestCase):