[Fix]: fix ndarray metainfo check in ConcatDataset

open-mmlab · Sep 1, 2023 · cdf1fda · cdf1fda
1 parent 8a7e80e
commit cdf1fda
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/mmengine/dataset/dataset_wrapper.py b/mmengine/dataset/dataset_wrapper.py
@@ -6,6 +6,7 @@
 from collections import defaultdict
 from typing import List, Sequence, Tuple, Union
 
+import numpy as np
 from torch.utils.data.dataset import ConcatDataset as _ConcatDataset
 
 from mmengine.logging import print_log
@@ -73,7 +74,14 @@ def __init__(self,
                     raise ValueError(
                         f'{key} does not in the meta information of '
                         f'the {i}-th dataset')
-                if self._metainfo[key] != dataset.metainfo[key]:
+                if isinstance(self._metainfo[key], np.ndarray) and isinstance(
+                        dataset.metainfo[key], np.ndarray):
+                    if not np.array_equal(self._metainfo[key],
+                                          dataset.metainfo[key]):
+                        raise ValueError(
+                            f'The meta information of the {i}-th dataset does'
+                            'not match meta information of the first dataset')
+                elif self._metainfo[key] != dataset.metainfo[key]:
                     raise ValueError(
                         f'The meta information of the {i}-th dataset does not '
                         'match meta information of the first dataset')