Use mmcv.FileClient to support different storage backends (#2712)

* use FileClient to replace mmcv.imread()

* typo fix

* rename file_client_cfg to file_client_args and bug fix

* add docstrings

* fix unittests

* fix unittests

mmdet/datasets/pipelines/loading.py

@@ -10,20 +10,48 @@
 
 @PIPELINES.register_module()
 class LoadImageFromFile(object):
+    """Load an image from file.
+
+    Required keys are "img_prefix" and "img_info" (a dict that must contain the
+    key "filename"). Added or updated keys are "filename", "img", "img_shape",
+    "ori_shape" (same as `img_shape`), "pad_shape" (same as `img_shape`),
+    "scale_factor" (1.0) and "img_norm_cfg" (means=0 and stds=1).
+
+    Args:
+        to_float32 (bool): Whether to convert the loaded image to a float32
+            numpy array. If set to False, the loaded image is an uint8 array.
+            Defaults to False.
+        color_type (str): The flag argument for :func:`mmcv.imfrombytes()`.
+            Defaults to 'color'.
+        file_client_args (dict): Arguments to instantiate a FileClient.
+            See :class:`mmcv.fileio.FileClient` for details.
+            Defaults to ``dict(backend='disk')``.
+    """
 
-    def __init__(self, to_float32=False, color_type='color'):
+    def __init__(self,
+                 to_float32=False,
+                 color_type='color',
+                 file_client_args=dict(backend='disk')):
         self.to_float32 = to_float32
         self.color_type = color_type
+        self.file_client_args = file_client_args.copy()
+        self.file_client = None
 
     def __call__(self, results):
+        if self.file_client is None:
+            self.file_client = mmcv.FileClient(**self.file_client_args)
+
         if results['img_prefix'] is not None:
             filename = osp.join(results['img_prefix'],
                                 results['img_info']['filename'])
         else:
             filename = results['img_info']['filename']
-        img = mmcv.imread(filename, self.color_type)
+
+        img_bytes = self.file_client.get(filename)
+        img = mmcv.imfrombytes(img_bytes, flag=self.color_type)
         if self.to_float32:
             img = img.astype(np.float32)
+
         results['filename'] = filename
         results['ori_filename'] = results['img_info']['filename']
         results['img'] = img
@@ -41,32 +69,63 @@ def __call__(self, results):
         return results
 
     def __repr__(self):
-        return f'{self.__class__.__name__}(to_float32={self.to_float32}, ' \
-            f"color_type='{self.color_type}')"
+        repr_str = (f'{self.__class__.__name__}('
+                    f'to_float32={self.to_float32}, '
+                    f"color_type='{self.color_type}', "
+                    f'file_client_args={self.file_client_args})')
+        return repr_str
 
 
 @PIPELINES.register_module()
 class LoadMultiChannelImageFromFiles(object):
-    """ Load multi channel images from a list of separate channel files.
-    Expects results['filename'] to be a list of filenames
+    """Load multi-channel images from a list of separate channel files.
+
+    Required keys are "img_prefix" and "img_info" (a dict that must contain the
+    key "filename", which is expected to be a list of filenames).
+    Added or updated keys are "filename", "img", "img_shape",
+    "ori_shape" (same as `img_shape`), "pad_shape" (same as `img_shape`),
+    "scale_factor" (1.0) and "img_norm_cfg" (means=0 and stds=1).
+
+    Args:
+        to_float32 (bool): Whether to convert the loaded image to a float32
+            numpy array. If set to False, the loaded image is an uint8 array.
+            Defaults to False.
+        color_type (str): The flag argument for :func:`mmcv.imfrombytes()`.
+            Defaults to 'color'.
+        file_client_args (dict): Arguments to instantiate a FileClient.
+            See :class:`mmcv.fileio.FileClient` for details.
+            Defaults to ``dict(backend='disk')``.
     """
 
-    def __init__(self, to_float32=False, color_type='unchanged'):
+    def __init__(self,
+                 to_float32=False,
+                 color_type='unchanged',
+                 file_client_args=dict(backend='disk')):
         self.to_float32 = to_float32
         self.color_type = color_type
+        self.file_client_args = file_client_args.copy()
+        self.file_client = None
 
     def __call__(self, results):
+        if self.file_client is None:
+            self.file_client = mmcv.FileClient(**self.file_client_args)
+
         if results['img_prefix'] is not None:
             filename = [
                 osp.join(results['img_prefix'], fname)
                 for fname in results['img_info']['filename']
             ]
         else:
             filename = results['img_info']['filename']
-        img = np.stack(
-            [mmcv.imread(name, self.color_type) for name in filename], axis=-1)
+
+        img = []
+        for name in filename:
+            img_bytes = self.file_client.get(name)
+            img.append(mmcv.imfrombytes(img_bytes, flag=self.color_type))
+        img = np.stack(img, axis=-1)
         if self.to_float32:
             img = img.astype(np.float32)
+
         results['filename'] = filename
         results['ori_filename'] = results['img_info']['filename']
         results['img'] = img
@@ -83,24 +142,47 @@ def __call__(self, results):
         return results
 
     def __repr__(self):
-        return f'{self.__class__.__name__}(to_float32={self.to_float32}, ' \
-            f"color_type='{self.color_type}')"
+        repr_str = (f'{self.__class__.__name__}('
+                    f'to_float32={self.to_float32}, '
+                    f"color_type='{self.color_type}', "
+                    f'file_client_args={self.file_client_args})')
+        return repr_str
 
 
 @PIPELINES.register_module()
 class LoadAnnotations(object):
+    """Load annotations.
+
+    Args:
+        with_bbox (bool): Whether to parse and load the bbox annotation.
+             Default: True.
+        with_label (bool): Whether to parse and load the label annotation.
+            Default: True.
+        with_mask (bool): Whether to parse and load the mask annotation.
+             Default: False.
+        with_seg (bool): Whether to parse and load the semantic segmentation
+            annotation. Default: False.
+        poly2mask (bool): Whether to convert the instance masks from polygons
+            to bitmaps. Default: True.
+        file_client_args (dict): Arguments to instantiate a FileClient.
+            See :class:`mmcv.fileio.FileClient` for details.
+            Defaults to ``dict(backend='disk')``.
+    """
 
     def __init__(self,
                  with_bbox=True,
                  with_label=True,
                  with_mask=False,
                  with_seg=False,
-                 poly2mask=True):
+                 poly2mask=True,
+                 file_client_args=dict(backend='disk')):
         self.with_bbox = with_bbox
         self.with_label = with_label
         self.with_mask = with_mask
         self.with_seg = with_seg
         self.poly2mask = poly2mask
+        self.file_client_args = file_client_args.copy()
+        self.file_client = None
 
     def _load_bboxes(self, results):
         ann_info = results['ann_info']
@@ -133,7 +215,7 @@ def _poly2mask(self, mask_ann, img_h, img_w):
         return mask
 
     def process_polygons(self, polygons):
-        """ Convert polygons to list of ndarray and filter invalid polygons.
+        """Convert polygons to list of ndarray and filter invalid polygons.
 
         Args:
             polygons (list[list]): polygons of one instance.
@@ -163,9 +245,14 @@ def _load_masks(self, results):
         return results
 
     def _load_semantic_seg(self, results):
-        results['gt_semantic_seg'] = mmcv.imread(
-            osp.join(results['seg_prefix'], results['ann_info']['seg_map']),
-            flag='unchanged').squeeze()
+        if self.file_client is None:
+            self.file_client = mmcv.FileClient(**self.file_client_args)
+
+        filename = osp.join(results['seg_prefix'],
+                            results['ann_info']['seg_map'])
+        img_bytes = self.file_client.get(filename)
+        results['gt_semantic_seg'] = mmcv.imfrombytes(
+            img_bytes, flag='unchanged').squeeze()
         results['seg_fields'].append('gt_semantic_seg')
         return results
 
@@ -188,6 +275,8 @@ def __repr__(self):
         repr_str += f'with_label={self.with_label}, '
         repr_str += f'with_mask={self.with_mask}, '
         repr_str += f'with_seg={self.with_seg})'
+        repr_str += f'poly2mask={self.poly2mask})'
+        repr_str += f'poly2mask={self.file_client_args})'
         return repr_str
 
 

tests/test_pipelines/test_loading.py

@@ -29,7 +29,8 @@ def test_load_img(self):
         np.testing.assert_equal(results['img_norm_cfg']['mean'],
                                 np.zeros(3, dtype=np.float32))
         assert repr(transform) == transform.__class__.__name__ + \
-            "(to_float32=False, color_type='color')"
+            "(to_float32=False, color_type='color', " + \
+            "file_client_args={'backend': 'disk'})"
 
         # no img_prefix
         results = dict(
@@ -78,4 +79,5 @@ def test_load_multi_channel_img(self):
         assert results['pad_shape'] == (288, 512, 3, 2)
         assert results['scale_factor'] == 1.0
         assert repr(transform) == transform.__class__.__name__ + \
-            "(to_float32=False, color_type='unchanged')"
+            "(to_float32=False, color_type='unchanged', " + \
+            "file_client_args={'backend': 'disk'})"