From 9e2815bbcf9560271f8a0adc4b76b6477c7f7e1e Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Tue, 26 Sep 2023 21:34:08 +0100 Subject: [PATCH] collate common metadict keys Signed-off-by: Wenqi Li --- monai/data/utils.py | 6 +++++- tests/test_list_data_collate.py | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/monai/data/utils.py b/monai/data/utils.py index 0880d44b64..8c5ae88289 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -452,7 +452,11 @@ def collate_meta_tensor(batch): elem_0 = first(batch) if isinstance(elem_0, MetaObj): collated = default_collate(batch) - collated.meta = default_collate([i.meta or TraceKeys.NONE for i in batch]) + meta_dicts = [i.meta or TraceKeys.NONE for i in batch] + common_ = set.intersection(*[set(d.keys()) for d in meta_dicts if isinstance(d, dict)]) + if common_: + meta_dicts = [{k: d[k] for k in common_} if isinstance(d, dict) else TraceKeys.NONE for d in meta_dicts] + collated.meta = default_collate(meta_dicts) collated.applied_operations = [i.applied_operations or TraceKeys.NONE for i in batch] collated.is_batch = True return collated diff --git a/tests/test_list_data_collate.py b/tests/test_list_data_collate.py index 226d9e1a55..9be61e3999 100644 --- a/tests/test_list_data_collate.py +++ b/tests/test_list_data_collate.py @@ -31,9 +31,13 @@ h = (np.array([19, 20, 21]), MetaTensor([22, 23, 24])) TEST_CASE_2 = [[[e, f], [g, h]], list, torch.Size([4, 3])] # dataset returns a list of tuple data +g_m = (np.array([13, 14, 15]), MetaTensor([16, 7, 18], meta={"key1": 0})) +h_m = (np.array([19, 20, 21]), MetaTensor([22, 23, 24], meta={"key2": 1})) +TEST_CASE_3 = [[[g_m], [h_m]], list, torch.Size([2, 3])] + class TestListDataCollate(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2]) + @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) def test_type_shape(self, input_data, expected_type, expected_shape): result = list_data_collate(input_data) self.assertIsInstance(result, expected_type)