Skip to content

5844 improve metatensor slicing #5845

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wyli merged 2 commits into from
Jan 13, 2023
Merged

5844 improve metatensor slicing #5845

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8250ca7
fixes slicing metatensor
wyli Jan 11, 2023
66e328a
Merge remote-tracking branch 'upstream' into 5844-slicing-metatensor
wyli Jan 13, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 12 additions & 4 deletions monai/data/meta_tensor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -235,11 +235,19 @@ def update_meta(rets: Sequence, func, args, kwargs) -> Sequence:
# respectively. Don't need to do anything with the metadata.
if batch_idx not in (slice(None, None, None), Ellipsis, None) and idx == 0:
ret_meta = decollate_batch(args[0], detach=False)[batch_idx]
if isinstance(ret_meta, list): # e.g. batch[0:2], re-collate
ret_meta = list_data_collate(ret_meta)
else: # e.g. `batch[0]` or `batch[0, 1]`, batch index is an integer
if isinstance(ret_meta, list) and ret_meta: # e.g. batch[0:2], re-collate
try:
ret_meta = list_data_collate(ret_meta)
except (TypeError, ValueError, RuntimeError, IndexError) as e:
raise ValueError(
"Inconsistent batched metadata dicts when slicing a batch of MetaTensors, "
"please convert it into a torch Tensor using `x.as_tensor()` or "
"a numpy array using `x.array`."
) from e
elif isinstance(ret_meta, MetaObj): # e.g. `batch[0]` or `batch[0, 1]`, batch_idx is int
ret_meta.is_batch = False
ret.__dict__ = ret_meta.__dict__.copy()
if hasattr(ret_meta, "__dict__"):
ret.__dict__ = ret_meta.__dict__.copy()
# `unbind` is used for `next(iter(batch))`. Also for `decollate_batch`.
# But we only want to split the batch if the `unbind` is along the 0th
# dimension.
Expand Down
7 changes: 7 additions & 0 deletions tests/test_meta_tensor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -406,6 +406,13 @@ def test_indexing(self):
for _d in d:
self.check_meta(_d, data)

def test_slicing(self):
x = MetaTensor(np.zeros((10, 3, 4)))
self.assertEqual(x[slice(4, 1)].shape[0], 0)
x.is_batch = True
with self.assertRaises(ValueError):
x[slice(0, 8)]

@parameterized.expand(DTYPES)
@SkipIfBeforePyTorchVersion((1, 8))
def test_decollate(self, dtype):
Expand Down