Image decoding optimization (#187)

* change image decoding method from cv2 to TurboJPEG
Signed-off-by: lixuanyi <lixuanyi@sensetime.com>

* add tests

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

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

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

* change per request
Signed-off-by: lizz <lizz@sensetime.com>

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

* baby-sitting
Signed-off-by: lizz <lizz@sensetime.com>
Co-authored-by: lizz <innerlee@users.noreply.github.com>

.travis.yml

@@ -5,7 +5,8 @@ language: python
 before_install:
   - sudo add-apt-repository -y ppa:mc3man/xerus-media
   - sudo apt-get update
-  - sudo apt-get install -y ffmpeg
+  - sudo apt-get install -y ffmpeg libjpeg-turbo*
+  - pip install -U git+git://github.com/lilohuang/PyTurboJPEG.git
 
 install:
   - rm -rf .eggs && pip install -e . codecov flake8 yapf isort mock

mmcv/image/__init__.py

 # Copyright (c) Open-MMLab. All rights reserved.
-from .io import imfrombytes, imread, imwrite
+from .io import imfrombytes, imread, imwrite, supported_backends, use_backend
 from .transforms import (bgr2gray, bgr2hls, bgr2hsv, bgr2rgb, gray2bgr,
                          gray2rgb, hls2bgr, hsv2bgr, imcrop, imdenormalize,
                          imflip, iminvert, imnormalize, impad,
@@ -11,5 +11,6 @@ __all__ = [
     'rgb2gray', 'gray2bgr', 'gray2rgb', 'bgr2rgb', 'rgb2bgr', 'bgr2hsv',
     'hsv2bgr', 'bgr2hls', 'hls2bgr', 'iminvert', 'imflip', 'imrotate',
     'imcrop', 'impad', 'impad_to_multiple', 'imnormalize', 'imdenormalize',
-    'imresize', 'imresize_like', 'imrescale'
+    'imresize', 'imresize_like', 'imrescale', 'use_backend',
+    'supported_backends'
 ]

mmcv/image/io.py

@@ -13,6 +13,13 @@ else:
     from cv2 import CV_LOAD_IMAGE_COLOR as IMREAD_COLOR
     from cv2 import CV_LOAD_IMAGE_GRAYSCALE as IMREAD_GRAYSCALE
     from cv2 import CV_LOAD_IMAGE_UNCHANGED as IMREAD_UNCHANGED
+try:
+    from turbojpeg import TJCS_RGB, TJPF_BGR, TJPF_GRAY, TurboJPEG
+except ImportError:
+    TJCS_RGB = TJPF_GRAY = TJPF_BGR = TurboJPEG = None
+
+jpeg = None
+supported_backends = ['cv2', 'turbojpeg']
 
 imread_flags = {
     'color': IMREAD_COLOR,
@@ -20,8 +27,43 @@ imread_flags = {
     'unchanged': IMREAD_UNCHANGED
 }
 
+imread_backend = 'cv2'
+
+
+def use_backend(backend):
+    """Select a backend for image decoding.
 
-def imread(img_or_path, flag='color'):
+    Args:
+        backend (str): The image decoding backend type. Options are `cv2` and
+            `turbojpeg` (see https://github.com/lilohuang/PyTurboJPEG).
+            `turbojpeg` is faster but it only supports `.jpeg` file format.
+    """
+    assert backend in supported_backends
+    global imread_backend
+    imread_backend = backend
+    if imread_backend == 'turbojpeg':
+        if TurboJPEG is None:
+            raise ValueError('`PyTurboJPEG` is not installed')
+        global jpeg
+        if jpeg is None:
+            jpeg = TurboJPEG()
+
+
+def _jpegflag(flag='color', channel_order='bgr'):
+    if flag == 'color':
+        if channel_order == 'bgr':
+            return TJPF_BGR
+        elif channel_order == 'rgb':
+            return TJCS_RGB
+        else:
+            raise ValueError('channel order must be "rgb" or "bgr"')
+    elif flag == 'grayscale':
+        return TJPF_GRAY
+    else:
+        raise ValueError('flag must be "color" or "grayscale"')
+
+
+def imread(img_or_path, flag='color', channel_order='bgr'):
     """Read an image.
 
     Args:
@@ -30,6 +72,8 @@ def imread(img_or_path, flag='color'):
             as is.
         flag (str): Flags specifying the color type of a loaded image,
             candidates are `color`, `grayscale` and `unchanged`.
+            Note that the `turbojpeg` backened does not support `unchanged`.
+        channel_order (str): Order of channel, candidates are `bgr` and `rgb`.
 
     Returns:
         ndarray: Loaded image array.
@@ -37,15 +81,26 @@ def imread(img_or_path, flag='color'):
     if isinstance(img_or_path, np.ndarray):
         return img_or_path
     elif is_str(img_or_path):
-        flag = imread_flags[flag] if is_str(flag) else flag
         check_file_exist(img_or_path,
                          'img file does not exist: {}'.format(img_or_path))
-        return cv2.imread(img_or_path, flag)
+        if imread_backend == 'turbojpeg':
+            with open(img_or_path, 'rb') as in_file:
+                img = jpeg.decode(in_file.read(),
+                                  _jpegflag(flag, channel_order))
+                if img.shape[-1] == 1:
+                    img = img[:, :, 0]
+            return img
+        else:
+            flag = imread_flags[flag] if is_str(flag) else flag
+            img = cv2.imread(img_or_path, flag)
+            if flag == IMREAD_COLOR and channel_order == 'rgb':
+                cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)
+            return img
     else:
         raise TypeError('"img" must be a numpy array or a filename')
 
 
-def imfrombytes(content, flag='color'):
+def imfrombytes(content, flag='color', channel_order='bgr'):
     """Read an image from bytes.
 
     Args:
@@ -55,10 +110,18 @@ def imfrombytes(content, flag='color'):
     Returns:
         ndarray: Loaded image array.
     """
-    img_np = np.frombuffer(content, np.uint8)
-    flag = imread_flags[flag] if is_str(flag) else flag
-    img = cv2.imdecode(img_np, flag)
-    return img
+    if imread_backend == 'turbojpeg':
+        img = jpeg.decode(content, _jpegflag(flag, channel_order))
+        if img.shape[-1] == 1:
+            img = img[:, :, 0]
+        return img
+    else:
+        img_np = np.frombuffer(content, np.uint8)
+        flag = imread_flags[flag] if is_str(flag) else flag
+        img = cv2.imdecode(img_np, flag)
+        if flag == IMREAD_COLOR and channel_order == 'rgb':
+            cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)
+        return img
 
 
 def imwrite(img, file_path, params=None, auto_mkdir=True):

tests/test_image.py

@@ -31,24 +31,80 @@ class TestImage(object):
         assert np.sum(diff <= 1) / float(area) > ratio_thr
 
     def test_imread(self):
-        img = mmcv.imread(self.img_path)
-        assert img.shape == (300, 400, 3)
-        img = mmcv.imread(self.img_path, 'grayscale')
-        assert img.shape == (300, 400)
-        img = mmcv.imread(self.gray_img_path)
-        assert img.shape == (300, 400, 3)
-        img = mmcv.imread(self.gray_img_path, 'unchanged')
-        assert img.shape == (300, 400)
-        img = mmcv.imread(img)
-        assert_array_equal(img, mmcv.imread(img))
+        # backend cv2
+        mmcv.use_backend('cv2')
+
+        img_cv2_color_bgr = mmcv.imread(self.img_path)
+        assert img_cv2_color_bgr.shape == (300, 400, 3)
+        img_cv2_color_rgb = mmcv.imread(self.img_path, channel_order='rgb')
+        assert img_cv2_color_rgb.shape == (300, 400, 3)
+        assert_array_equal(img_cv2_color_rgb[:, :, ::-1], img_cv2_color_bgr)
+        img_cv2_grayscale1 = mmcv.imread(self.img_path, 'grayscale')
+        assert img_cv2_grayscale1.shape == (300, 400)
+        img_cv2_grayscale2 = mmcv.imread(self.gray_img_path)
+        assert img_cv2_grayscale2.shape == (300, 400, 3)
+        img_cv2_unchanged = mmcv.imread(self.gray_img_path, 'unchanged')
+        assert img_cv2_unchanged.shape == (300, 400)
+        img_cv2_unchanged = mmcv.imread(img_cv2_unchanged)
+        assert_array_equal(img_cv2_unchanged, mmcv.imread(img_cv2_unchanged))
         with pytest.raises(TypeError):
             mmcv.imread(1)
 
+        # backend turbojpeg
+        mmcv.use_backend('turbojpeg')
+
+        img_turbojpeg_color_bgr = mmcv.imread(self.img_path)
+        assert img_turbojpeg_color_bgr.shape == (300, 400, 3)
+        assert_array_equal(img_turbojpeg_color_bgr, img_cv2_color_bgr)
+
+        img_turbojpeg_color_rgb = mmcv.imread(
+            self.img_path, channel_order='rgb')
+        assert img_turbojpeg_color_rgb.shape == (300, 400, 3)
+        assert_array_equal(img_turbojpeg_color_rgb, img_cv2_color_rgb)
+
+        with pytest.raises(ValueError):
+            mmcv.imread(self.img_path, channel_order='unsupport_order')
+
+        img_turbojpeg_grayscale1 = mmcv.imread(self.img_path, flag='grayscale')
+        assert img_turbojpeg_grayscale1.shape == (300, 400)
+        assert_array_equal(img_turbojpeg_grayscale1, img_cv2_grayscale1)
+
+        img_turbojpeg_grayscale2 = mmcv.imread(self.gray_img_path)
+        assert img_turbojpeg_grayscale2.shape == (300, 400, 3)
+        assert_array_equal(img_turbojpeg_grayscale2, img_cv2_grayscale2)
+
+        img_turbojpeg_grayscale2 = mmcv.imread(img_turbojpeg_grayscale2)
+        assert_array_equal(img_turbojpeg_grayscale2,
+                           mmcv.imread(img_turbojpeg_grayscale2))
+
+        with pytest.raises(ValueError):
+            mmcv.imread(self.gray_img_path, 'unchanged')
+
+        with pytest.raises(TypeError):
+            mmcv.imread(1)
+
+        with pytest.raises(AssertionError):
+            mmcv.use_backend('unsupport_backend')
+
+        mmcv.use_backend('cv2')
+
     def test_imfrombytes(self):
+        # backend cv2
+        mmcv.use_backend('cv2')
         with open(self.img_path, 'rb') as f:
             img_bytes = f.read()
-        img = mmcv.imfrombytes(img_bytes)
-        assert img.shape == (300, 400, 3)
+        img_cv2 = mmcv.imfrombytes(img_bytes)
+        assert img_cv2.shape == (300, 400, 3)
+
+        # backend turbojpeg
+        mmcv.use_backend('turbojpeg')
+        with open(self.img_path, 'rb') as f:
+            img_bytes = f.read()
+        img_turbojpeg = mmcv.imfrombytes(img_bytes)
+        assert img_turbojpeg.shape == (300, 400, 3)
+        assert_array_equal(img_cv2, img_turbojpeg)
+
+        mmcv.use_backend('cv2')
 
     def test_imwrite(self):
         img = mmcv.imread(self.img_path)