diff --git a/mmcv/image/__init__.py b/mmcv/image/__init__.py index de2fda4791..60e9d87638 100644 --- a/mmcv/image/__init__.py +++ b/mmcv/image/__init__.py @@ -8,9 +8,9 @@ from .io import imfrombytes, imread, imwrite, supported_backends, use_backend from .misc import tensor2imgs from .photometric import (adjust_brightness, adjust_color, adjust_contrast, - adjust_sharpness, clahe, imdenormalize, imequalize, - iminvert, imnormalize, imnormalize_, lut_transform, - posterize, solarize) + adjust_sharpness, auto_contrast, clahe, + imdenormalize, imequalize, iminvert, imnormalize, + imnormalize_, lut_transform, posterize, solarize) __all__ = [ 'bgr2gray', 'bgr2hls', 'bgr2hsv', 'bgr2rgb', 'gray2bgr', 'gray2rgb', @@ -22,5 +22,5 @@ 'rgb2ycbcr', 'bgr2ycbcr', 'ycbcr2rgb', 'ycbcr2bgr', 'tensor2imgs', 'imshear', 'imtranslate', 'adjust_color', 'imequalize', 'adjust_brightness', 'adjust_contrast', 'lut_transform', 'clahe', - 'adjust_sharpness' + 'adjust_sharpness', 'auto_contrast' ] diff --git a/mmcv/image/photometric.py b/mmcv/image/photometric.py index 6dbd87bd36..60be233f71 100644 --- a/mmcv/image/photometric.py +++ b/mmcv/image/photometric.py @@ -232,6 +232,62 @@ def adjust_contrast(img, factor=1.): return contrasted_img.astype(img.dtype) +def auto_contrast(img, cutoff=0): + """Auto adjust image contrast. + + This function maximize (normalize) image contrast by first removing cutoff + percent of the lightest and darkest pixels from the histogram and remapping + the image so that the darkest pixel becomes black (0), and the lightest + becomes white (255). + + Args: + img (ndarray): Image to be contrasted. BGR order. + cutoff (int | float | tuple): The cutoff percent of the lightest and + darkest pixels to be removed. If given as tuple, it shall be + (low, high). Otherwise, the single value will be used for both. + Defaults to 0. + + Returns: + ndarray: The contrasted image. + """ + + def _auto_contrast_channel(im, c, cutoff): + im = im[:, :, c] + # Compute the histogram of the image channel. + histo = np.histogram(im, 256, (0, 255))[0] + # Remove cut-off percent pixels from histo + histo_sum = np.cumsum(histo) + cut_low = histo_sum[-1] * cutoff[0] // 100 + cut_high = histo_sum[-1] - histo_sum[-1] * cutoff[1] // 100 + histo_sum = np.clip(histo_sum, cut_low, cut_high) - cut_low + histo = np.concatenate([[histo_sum[0]], np.diff(histo_sum)], 0) + + # Compute mapping + low, high = np.nonzero(histo)[0][0], np.nonzero(histo)[0][-1] + # If all the values have been cut off, return the origin img + if low >= high: + return im + scale = 255.0 / (high - low) + offset = -low * scale + lut = np.array(range(256)) + lut = lut * scale + offset + lut = np.clip(lut, 0, 255) + return lut[im] + + if isinstance(cutoff, (int, float)): + cutoff = (cutoff, cutoff) + else: + assert isinstance(cutoff, tuple), 'cutoff must be of type int, ' \ + f'float or tuple, but got {type(cutoff)} instead.' + # Auto adjusts contrast for each channel independently and then stacks + # the result. + s1 = _auto_contrast_channel(img, 0, cutoff) + s2 = _auto_contrast_channel(img, 1, cutoff) + s3 = _auto_contrast_channel(img, 2, cutoff) + contrasted_img = np.stack([s1, s2, s3], axis=-1) + return contrasted_img.astype(img.dtype) + + def adjust_sharpness(img, factor=1., kernel=None): """Adjust image sharpness. diff --git a/tests/test_image/test_photometric.py b/tests/test_image/test_photometric.py index 9c2676eddb..bee209d432 100644 --- a/tests/test_image/test_photometric.py +++ b/tests/test_image/test_photometric.py @@ -202,6 +202,50 @@ def _adjust_contrast(img, factor): rtol=0, atol=1) + def test_auto_contrast(self, nb_rand_test=100): + + def _auto_contrast(img, cutoff=0): + from PIL.ImageOps import autocontrast + from PIL import Image + # Image.fromarray defaultly supports RGB, not BGR. + # convert from BGR to RGB + img = Image.fromarray(img[..., ::-1], mode='RGB') + contrasted_img = autocontrast(img, cutoff) + # convert from RGB to BGR + return np.asarray(contrasted_img)[..., ::-1] + + img = np.array([[0, 128, 255], [1, 127, 254], [2, 129, 253]], + dtype=np.uint8) + img = np.stack([img, img, img], axis=-1) + + # test case without cut-off + assert_array_equal(mmcv.auto_contrast(img), _auto_contrast(img)) + # test case with cut-off as int + assert_array_equal( + mmcv.auto_contrast(img, 10), _auto_contrast(img, 10)) + # test case with cut-off as float + assert_array_equal( + mmcv.auto_contrast(img, 12.5), _auto_contrast(img, 12.5)) + # test case with cut-off as tuple + assert_array_equal( + mmcv.auto_contrast(img, (10, 10)), _auto_contrast(img, 10)) + # test case with cut-off with sum over 100 + assert_array_equal( + mmcv.auto_contrast(img, 60), _auto_contrast(img, 60)) + + # test auto_contrast with randomly sampled images and factors. + for _ in range(nb_rand_test): + img = np.clip( + np.random.uniform(0, 1, (1200, 1000, 3)) * 260, 0, + 255).astype(np.uint8) + # cut-offs are not set as tuple since in `build.yml`, pillow 6.2.2 + # is installed, which does not support setting low cut-off and high + # cut-off differently. + # With pillow above 8.0.0, cutoff can be set as tuple + cutoff = np.random.rand() * 100 + assert_array_equal( + mmcv.auto_contrast(img, cutoff), _auto_contrast(img, cutoff)) + def test_adjust_sharpness(self, nb_rand_test=100): def _adjust_sharpness(img, factor):