Class Contrast

tensorpack/dataflow/imgaug/imgproc.py:99–133

Apply ``x = (x - mean) * contrast_factor + mean`` to each channel.

class Contrast(PhotometricAugmentor):
    """
    Apply ``x = (x - mean) * contrast_factor + mean`` to each channel.
    """

    def __init__(self, factor_range, rgb=None, clip=True):
        """
        Args:
            factor_range (list or tuple): an interval to randomly sample the `contrast_factor`.
            rgb (bool or None): if None, use the mean per-channel.
            clip (bool): clip to [0, 255] even when data type is not uint8.
        """
        super(Contrast, self).__init__()
        self._init(locals())

    def _get_augment_params(self, _):
        return self._rand_range(*self.factor_range)

    def _augment(self, img, r):
        old_dtype = img.dtype

        if img.ndim == 3:
            if self.rgb is not None:
                m = cv2.COLOR_RGB2GRAY if self.rgb else cv2.COLOR_BGR2GRAY
                grey = cv2.cvtColor(img.astype('float32'), m)
                mean = np.mean(grey)
            else:
                mean = np.mean(img, axis=(0, 1), keepdims=True)
        else:
            mean = np.mean(img)

        img = img * r + mean * (1 - r)
        if self.clip or old_dtype == np.uint8:
            img = np.clip(img, 0, 255)
        return img.astype(old_dtype)


class MeanVarianceNormalize(PhotometricAugmentor):