add inplace normalization (#192)

* add inplace normalization
Signed-off-by: lixuanyi <lixuanyi@sensetime.com>

* fix
Signed-off-by: lixuanyi <lixuanyi@sensetime.com>

* add dtype check

* modify inplace norm
Signed-off-by: lixuanyi <lixuanyi@sensetime.com>

* fix
Signed-off-by: lixuanyi <lixuanyi@sensetime.com>

* fix
Signed-off-by: lixuanyi <lixuanyi@sensetime.com>

* fix line

mmcv/image/__init__.py

@@ -2,15 +2,16 @@
 from .io import imfrombytes, imread, imwrite, supported_backends, use_backend
 from .transforms import (bgr2gray, bgr2hls, bgr2hsv, bgr2rgb, gray2bgr,
                          gray2rgb, hls2bgr, hsv2bgr, imcrop, imdenormalize,
-                         imflip, imflip_, iminvert, imnormalize, impad,
-                         impad_to_multiple, imrescale, imresize, imresize_like,
-                         imrotate, posterize, rgb2bgr, rgb2gray, solarize)
+                         imflip, imflip_, iminvert, imnormalize, imnormalize_,
+                         impad, impad_to_multiple, imrescale, imresize,
+                         imresize_like, imrotate, posterize, rgb2bgr, rgb2gray,
+                         solarize)
 
 __all__ = [
     'solarize', 'posterize', 'imread', 'imwrite', 'imfrombytes', 'bgr2gray',
     'rgb2gray', 'gray2bgr', 'gray2rgb', 'bgr2rgb', 'rgb2bgr', 'bgr2hsv',
     'hsv2bgr', 'bgr2hls', 'hls2bgr', 'iminvert', 'imflip', 'imflip_',
     'imrotate', 'imcrop', 'impad', 'impad_to_multiple', 'imnormalize',
-    'imdenormalize', 'imresize', 'imresize_like', 'imrescale', 'use_backend',
-    'supported_backends'
+    'imnormalize_', 'imdenormalize', 'imresize', 'imresize_like', 'imrescale',
+    'use_backend', 'supported_backends'
 ]

mmcv/image/transforms/__init__.py

@@ -4,13 +4,13 @@ from .colorspace import (bgr2gray, bgr2hls, bgr2hsv, bgr2rgb, gray2bgr,
                          rgb2bgr, rgb2gray, solarize)
 from .geometry import (imcrop, imflip, imflip_, impad, impad_to_multiple,
                        imrotate)
-from .normalize import imdenormalize, imnormalize
+from .normalize import imdenormalize, imnormalize, imnormalize_
 from .resize import imrescale, imresize, imresize_like
 
 __all__ = [
     'solarize', 'posterize', 'bgr2gray', 'rgb2gray', 'gray2bgr', 'gray2rgb',
     'bgr2rgb', 'rgb2bgr', 'bgr2hsv', 'hsv2bgr', 'bgr2hls', 'hls2bgr',
     'iminvert', 'imflip', 'imflip_', 'imrotate', 'imcrop', 'impad',
-    'impad_to_multiple', 'imnormalize', 'imdenormalize', 'imresize',
-    'imresize_like', 'imrescale'
+    'impad_to_multiple', 'imnormalize', 'imnormalize_', 'imdenormalize',
+    'imresize', 'imresize_like', 'imrescale'
 ]

mmcv/image/transforms/normalize.py

@@ -4,7 +4,35 @@ import numpy as np
 
 
 def imnormalize(img, mean, std, to_rgb=True):
+    """Normalize an image with mean and std.
+
+    Args:
+        img (ndarray): Image to be normalized.
+        mean (ndarray): The mean to be used for normalize.
+        std (ndarray): The std to be used for normalize.
+        to_rgb (bool): Whether to convert to rgb.
+
+    Returns:
+        ndarray: The normalized image.
+    """
     img = np.float32(img) if img.dtype != np.float32 else img.copy()
+    return imnormalize_(img, mean, std, to_rgb)
+
+
+def imnormalize_(img, mean, std, to_rgb=True):
+    """Inplace normalize an image with mean and std.
+
+    Args:
+        img (ndarray): Image to be normalized.
+        mean (ndarray): The mean to be used for normalize.
+        std (ndarray): The std to be used for normalize.
+        to_rgb (bool): Whether to convert to rgb.
+
+    Returns:
+        ndarray: The normalized image.
+    """
+    # cv2 inplace normalization does not accept uint8
+    assert img.dtype != np.uint8
     mean = np.float64(mean.reshape(1, -1))
     stdinv = 1 / np.float64(std.reshape(1, -1))
     if to_rgb:

tests/test_image.py

@@ -119,8 +119,22 @@ class TestImage(object):
         baseline = (rgbimg - self.mean) / self.std
         img = mmcv.imnormalize(self.img, self.mean, self.std)
         assert np.allclose(img, baseline)
+        assert id(img) != id(self.img)
         img = mmcv.imnormalize(rgbimg, self.mean, self.std, to_rgb=False)
         assert np.allclose(img, baseline)
+        assert id(img) != id(rgbimg)
+
+    def test_imnormalize_(self):
+        img_for_normalize = np.float32(self.img.copy())
+        rgbimg_for_normalize = np.float32(self.img[:, :, ::-1].copy())
+        baseline = (rgbimg_for_normalize - self.mean) / self.std
+        img = mmcv.imnormalize_(img_for_normalize, self.mean, self.std)
+        assert np.allclose(img_for_normalize, baseline)
+        assert id(img) == id(img_for_normalize)
+        img = mmcv.imnormalize_(
+            rgbimg_for_normalize, self.mean, self.std, to_rgb=False)
+        assert np.allclose(img, baseline)
+        assert id(img) == id(rgbimg_for_normalize)
 
     def test_imdenormalize(self):
         normimg = (self.img[:, :, ::-1] - self.mean) / self.std