Even more tests

Signed-off-by: Kamil Tokarski <[email protected]>

dali/python/nvidia/dali/auto_aug/augmentations.py

@@ -109,9 +109,28 @@ def brightness(sample, parameter):
 
 @augmentation(mag_range=(0, 0.9), randomly_negate=True, as_param=shift_enhance_range)
 def contrast(sample, parameter):
+    """
+    This variant is not PIL-conformant as it scales the contrast around the center of
+    type range rather than a channel-weighted mean as PIL does. See `contrast_mean_centered`.
+    """
     return fn.contrast(sample, contrast=parameter)
 
 
+@augmentation(mag_range=(0, 0.9), randomly_negate=True, as_param=shift_enhance_range)
+def contrast_mean_centered(sample, parameter):
+    """
+    This variant follows PIL implementation of Contrast enhancement, which does not use
+    the middle of the `dtype` range as the contrast center, but a channel-weighted mean.
+    NOTE: Currently, it is not possible to run this variant with `sample` that resides
+    in GPU memory due to reduction values residing in GPU and the limitation that named arguments
+    (i.e. contrast_center) must reside in CPU memory.
+    """
+    mean = fn.reductions.mean(sample, axes=[0, 1])
+    rgb_weights = types.Constant(np.array([0.299, 0.587, 0.114], dtype=np.float32))
+    center = fn.reductions.sum(mean * rgb_weights)
+    return fn.contrast(sample, contrast=parameter, contrast_center=center)
+
+
 @augmentation(mag_range=(0, 0.9), randomly_negate=True, as_param=shift_enhance_range)
 def color(sample, parameter):
     return fn.saturation(sample, saturation=parameter)
@@ -185,6 +204,10 @@ def invert(sample, _):
 
 @augmentation
 def equalize(sample, _):
+    """
+    DALI's equalize is open-cv conformant; the PIL impl uses slightly different formula when
+    scaling histogram's cumulative sum to create lookup table
+    """
     return fn.experimental.equalize(sample)
 
 

dali/test/python/auto_aug/test_augmentations.py

@@ -179,6 +179,21 @@ def contrast_ref(input, contrast):
                              batch_size=batch_size, max_allowed_error=1, dev=dev, params=magnitudes)
 
 
+@params(("cpu", ))
+def test_contrast_mean_centered(dev):
+
+    def contrast_ref(img, magnitude):
+        return ImageEnhance.Contrast(img).enhance(magnitude)
+
+    batch_size = 16
+    data_source = get_images(dev)
+    contrast = a.contrast_mean_centered.augmentation(mag_range=(0.1, 1), randomly_negate=False,
+                                                     as_param=None)
+    magnitudes = contrast._get_magnitudes(batch_size)
+    compare_against_baseline(contrast, pil_baseline(contrast_ref), data_source,
+                             batch_size=batch_size, max_allowed_error=1, dev=dev, params=magnitudes)
+
+
 @params(("cpu", ), ("gpu", ))
 def test_color(dev):
     max_allowed_error = 2