[Enhancement] Allow to set channel_order in LoadImageFromFile (#7258)

* allow to set channel_order when loading images * fix lint * fix unit test * fix lint
open-mmlab · Mar 16, 2022 · f5d3b7c · f5d3b7c
1 parent 075d6f8
commit f5d3b7c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/mmdet/datasets/pipelines/loading.py b/mmdet/datasets/pipelines/loading.py
@@ -37,9 +37,11 @@ class LoadImageFromFile:
     def __init__(self,
                  to_float32=False,
                  color_type='color',
+                 channel_order='bgr',
                  file_client_args=dict(backend='disk')):
         self.to_float32 = to_float32
         self.color_type = color_type
+        self.channel_order = channel_order
         self.file_client_args = file_client_args.copy()
         self.file_client = None
 
@@ -63,7 +65,8 @@ def __call__(self, results):
             filename = results['img_info']['filename']
 
         img_bytes = self.file_client.get(filename)
-        img = mmcv.imfrombytes(img_bytes, flag=self.color_type)
+        img = mmcv.imfrombytes(
+            img_bytes, flag=self.color_type, channel_order=self.channel_order)
         if self.to_float32:
             img = img.astype(np.float32)
 
@@ -79,6 +82,7 @@ def __repr__(self):
         repr_str = (f'{self.__class__.__name__}('
                     f'to_float32={self.to_float32}, '
                     f"color_type='{self.color_type}', "
+                    f"channel_order='{self.channel_order}', "
                     f'file_client_args={self.file_client_args})')
         return repr_str
 

diff --git a/tests/test_data/test_pipelines/test_loading.py b/tests/test_data/test_pipelines/test_loading.py
@@ -27,7 +27,7 @@ def test_load_img(self):
         assert results['img_shape'] == (288, 512, 3)
         assert results['ori_shape'] == (288, 512, 3)
         assert repr(transform) == transform.__class__.__name__ + \
-            "(to_float32=False, color_type='color', " + \
+            "(to_float32=False, color_type='color', channel_order='bgr', " + \
             "file_client_args={'backend': 'disk'})"
 
         # no img_prefix