transforms/loading: support reading numpy files (.npy, .npz)

open-mmlab · Oct 5, 2023 · f6865bd · f6865bd
1 parent f2439cd
commit f6865bd
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 10 deletions.
diff --git a/docs/en/understand_mmcv/data_process.md b/docs/en/understand_mmcv/data_process.md
@@ -2,11 +2,11 @@
 
 ### Image
 
-This module provides some image processing methods, which requires `opencv` to be installed first.
+This module provides some image processing methods, which requires `opencv` and `numpy` to be installed first.
 
 #### Read/Write/Show
 
-To read or write images files, use `imread` or `imwrite`.
+To read or write images files, use `imread` or `imwrite` with `opencv` and `load` or `save` with `numpy`.
 
 ```python
 import mmcv
@@ -17,6 +17,13 @@ img_ = mmcv.imread(img)  # nothing will happen, img_ = img
 mmcv.imwrite(img, 'out.jpg')
 ```
 
+```python
+import numpy as np
+
+img = np.load('test.npy')
+np.save('test.npy', img)
+```
+
 To read images from bytes
 
 ```python

diff --git a/mmcv/transforms/loading.py b/mmcv/transforms/loading.py
@@ -1,5 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import warnings
+from pathlib import Path
 from typing import Optional
 
 import mmengine.fileio as fileio
@@ -90,21 +91,28 @@ def transform(self, results: dict) -> Optional[dict]:
         """
 
         filename = results['img_path']
+
         try:
-            if self.file_client_args is not None:
-                file_client = fileio.FileClient.infer_client(
-                    self.file_client_args, filename)
-                img_bytes = file_client.get(filename)
+            if Path(filename).suffix in ['.npy', '.npz']:
+                img = np.load(filename)
             else:
-                img_bytes = fileio.get(
-                    filename, backend_args=self.backend_args)
-            img = mmcv.imfrombytes(
-                img_bytes, flag=self.color_type, backend=self.imdecode_backend)
+                if self.file_client_args is not None:
+                    file_client = fileio.FileClient.infer_client(
+                        self.file_client_args, filename)
+                    img_bytes = file_client.get(filename)
+                else:
+                    img_bytes = fileio.get(
+                        filename, backend_args=self.backend_args)
+                img = mmcv.imfrombytes(
+                    img_bytes,
+                    flag=self.color_type,
+                    backend=self.imdecode_backend)
         except Exception as e:
             if self.ignore_empty:
                 return None
             else:
                 raise e
+
         # in some cases, images are not read successfully, the img would be
         # `None`, refer to https://github.com/open-mmlab/mmpretrain/issues/1427
         assert img is not None, f'failed to load image: {filename}'