From 366c628a85daf71fa26f652b402277c90d9f7d6e Mon Sep 17 00:00:00 2001 From: Ma Zerun Date: Wed, 14 Jul 2021 21:03:06 +0800 Subject: [PATCH] [Fix] Fix potential interger overflow in `imequalize`. (#1198) * Fix potential interger overflow in `imequalize`. * Modify imequalize unit test image size to generate potential integer overflow. --- mmcv/image/photometric.py | 2 ++ tests/test_image/test_photometric.py | 5 ++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/mmcv/image/photometric.py b/mmcv/image/photometric.py index c43c33dd99..a26cf5b597 100644 --- a/mmcv/image/photometric.py +++ b/mmcv/image/photometric.py @@ -157,6 +157,8 @@ def _scale_channel(im, c): lut = (np.cumsum(histo) + (step // 2)) // step # Shift lut, prepending with 0. lut = np.concatenate([[0], lut[:-1]], 0) + # handle potential integer overflow + lut[lut > 255] = 255 # If step is zero, return the original image. # Otherwise, index from lut. return np.where(np.equal(step, 0), im, lut[im]) diff --git a/tests/test_image/test_photometric.py b/tests/test_image/test_photometric.py index 3f43f5ca69..cff805688e 100644 --- a/tests/test_image/test_photometric.py +++ b/tests/test_image/test_photometric.py @@ -131,9 +131,8 @@ def _imequalize(img): # test equalize with randomly sampled image. for _ in range(nb_rand_test): - img = np.clip( - np.random.normal(0, 1, (1000, 1200, 3)) * 260, 0, - 255).astype(np.uint8) + img = np.clip(np.random.normal(0, 1, (256, 256, 3)) * 260, 0, + 255).astype(np.uint8) equalized_img = mmcv.imequalize(img) assert_array_equal(equalized_img, _imequalize(img))