Use cv2 to normalize img (#176)

* Use cv2 to normalize img

* multipy faster than divide

* also do imdenormalize

* support more channels

* even faster
Signed-off-by: lizz <lizz@sensetime.com>

* add test
Signed-off-by: lizz <lizz@sensetime.com>

* even even faster
Signed-off-by: lizz <lizz@sensetime.com>

* fffaster
Signed-off-by: lizz <lizz@sensetime.com>

* okay
Signed-off-by: lizz <lizz@sensetime.com>

* guard against implace modification
Signed-off-by: lizz <lizz@sensetime.com>

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

* tidy
Signed-off-by: lizz <lizz@sensetime.com>

* clean
Signed-off-by: lizz <lizz@sensetime.com>

* Update test_image.py

* Update test_image.py

* Update test_image.py

* Update test_image.py

* Update test_image.py

* Update test_image.py

mmcv/image/transforms/normalize.py

 # Copyright (c) Open-MMLab. All rights reserved.
+import cv2
 import numpy as np
 
-from .colorspace import bgr2rgb, rgb2bgr
-
 
 def imnormalize(img, mean, std, to_rgb=True):
-    img = img.astype(np.float32)
+    img = np.float32(img) if img.dtype != np.float32 else img.copy()
+    mean = np.float64(mean.reshape(1, -1))
+    stdinv = 1 / np.float64(std.reshape(1, -1))
     if to_rgb:
-        img = bgr2rgb(img)
-    return (img - mean) / std
+        cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)  # inplace
+    cv2.subtract(img, mean, img)  # inplace
+    cv2.multiply(img, stdinv, img)  # inplace
+    return img
 
 
 def imdenormalize(img, mean, std, to_bgr=True):
-    img = (img * std) + mean
+    assert img.dtype != np.uint8
+    mean = mean.reshape(1, -1).astype(np.float64)
+    std = std.reshape(1, -1).astype(np.float64)
+    img = cv2.multiply(img, std)  # make a copy
+    cv2.add(img, mean, img)  # inplace
     if to_bgr:
-        img = rgb2bgr(img)
+        cv2.cvtColor(img, cv2.COLOR_RGB2BGR, img)  # inplace
     return img

tests/test_image.py

@@ -20,6 +20,8 @@ class TestImage(object):
         cls.gray_img_path = osp.join(
             osp.dirname(__file__), 'data/grayscale.jpg')
         cls.img = cv2.imread(cls.img_path)
+        cls.mean = np.float32(np.array([123.675, 116.28, 103.53]))
+        cls.std = np.float32(np.array([58.395, 57.12, 57.375]))
 
     def assert_img_equal(self, img, ref_img, ratio_thr=0.999):
         assert img.shape == ref_img.shape
@@ -56,6 +58,23 @@ class TestImage(object):
         os.remove(out_file)
         self.assert_img_equal(img, rewrite_img)
 
+    def test_imnormalize(self):
+        rgbimg = self.img[:, :, ::-1]
+        baseline = (rgbimg - self.mean) / self.std
+        img = mmcv.imnormalize(self.img, self.mean, self.std)
+        assert np.allclose(img, baseline)
+        img = mmcv.imnormalize(rgbimg, self.mean, self.std, to_rgb=False)
+        assert np.allclose(img, baseline)
+
+    def test_imdenormalize(self):
+        normimg = (self.img[:, :, ::-1] - self.mean) / self.std
+        rgbbaseline = (normimg * self.std + self.mean)
+        bgrbaseline = rgbbaseline[:, :, ::-1]
+        img = mmcv.imdenormalize(normimg, self.mean, self.std)
+        assert np.allclose(img, bgrbaseline)
+        img = mmcv.imdenormalize(normimg, self.mean, self.std, to_bgr=False)
+        assert np.allclose(img, rgbbaseline)
+
     def test_bgr2gray(self):
         in_img = np.random.rand(10, 10, 3).astype(np.float32)
         out_img = mmcv.bgr2gray(in_img)