collate common metadict keys

Signed-off-by: Wenqi Li <[email protected]>

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

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)