Skip to content

Commit

Permalink
7042 collate common meta dictionary keys (#7054)
Browse files Browse the repository at this point in the history
Fixes #7042

- handling inconsistent meta keys when collating metatensor


### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [x] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [x] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

Signed-off-by: Wenqi Li <[email protected]>
  • Loading branch information
wyli authored Sep 27, 2023
1 parent f214b27 commit a29ab04
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 2 deletions.
6 changes: 5 additions & 1 deletion monai/data/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 5 additions & 1 deletion tests/test_list_data_collate.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit a29ab04

Please sign in to comment.