Repository: Project-MONAI/MONAI

Issue:
metatensor slicing should decollate applied operations
**Describe the bug**
```py
from monai.data import MetaTensor, list_data_collate
from monai.transforms import Resize

a = MetaTensor([[1.0, 2.0]])
a = Resize((1, 2))(a)
b = list_data_collate([a, a])
print(b[0])  # decollate, slicing
```
output:
```
tensor([[[1.5000, 1.5000]]])
MetaData
        affine: tensor([[ 2.0000,  0.0000,  0.0000,  0.5000],
        [ 0.0000,  0.5000,  0.0000, -0.2500],
        [ 0.0000,  0.0000,  1.0000,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  1.0000]], dtype=torch.float64)

Applied operations
[ [ { class: 'Resize',
      extra_info: {'align_corners': none, 'mode': area, 'new_dim': 1},
      id: 140688499272720,
      orig_size: (2, 1)}],
  [ { class: 'Resize',
      extra_info: {'align_corners': none, 'mode': area, 'new_dim': 1},
      id: 140688499272720,
      orig_size: (2, 1)}]]
Is batch?: False
```

the expected output is a metatensor with a single element in `applied_operations`


Discussion:

Resolving patch:
diff --git a/monai/data/meta_tensor.py b/monai/data/meta_tensor.py
--- a/monai/data/meta_tensor.py
+++ b/monai/data/meta_tensor.py
@@ -215,20 +215,13 @@ def update_meta(rets: Sequence, func, args, kwargs) -> Sequence:
                         # if using e.g., `batch[:, -1]` or `batch[..., -1]`, then the
                         # first element will be `slice(None, None, None)` and `Ellipsis`,
                         # respectively. Don't need to do anything with the metadata.
-                        if batch_idx not in (slice(None, None, None), Ellipsis, None):
-                            # only decollate metadata once
-                            if metas is None:
-                                metas = decollate_batch(ret.meta)
-                            meta = metas[batch_idx]
-                            # if using e.g., `batch[0:2]`, then `is_batch` should still be
-                            # `True`. Also re-collate the remaining elements.
-                            if isinstance(meta, list):
-                                ret.meta = list_data_collate(meta)
-                            # if using e.g., `batch[0]` or `batch[0, 1]`, then return single
-                            # element from batch, and set `is_batch` to `False`.
-                            else:
-                                ret.meta = meta
-                                ret.is_batch = False
+                        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
+                                ret_meta.is_batch = False
+                            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.
@@ -241,8 +234,8 @@ def update_meta(rets: Sequence, func, args, kwargs) -> Sequence:
                             dim = 0
                         if dim == 0:
                             if metas is None:
-                                metas = decollate_batch(ret.meta)
-                            ret.meta = metas[idx]
+                                metas = decollate_batch(args[0], detach=False)
+                            ret.__dict__ = metas[idx].__dict__.copy()
                             ret.is_batch = False
 
             out.append(ret)
@@ -330,7 +323,7 @@ def as_tensor(self) -> torch.Tensor:
         """
         return self.as_subclass(torch.Tensor)  # type: ignore
 
-    def get_array(self, output_type=np.ndarray, dtype=None, *_args, **_kwargs):
+    def get_array(self, output_type=np.ndarray, dtype=None, device=None, *_args, **_kwargs):
         """
         Returns a new array in `output_type`, the array shares the same underlying storage when the output is a
         numpy array. Changes to self tensor will be reflected in the ndarray and vice versa.
@@ -340,10 +333,11 @@ def get_array(self, output_type=np.ndarray, dtype=None, *_args, **_kwargs):
             dtype: dtype of output data. Converted to correct library type (e.g.,
                 `np.float32` is converted to `torch.float32` if output type is `torch.Tensor`).
                 If left blank, it remains unchanged.
+            device: if the output is a `torch.Tensor`, select device (if `None`, unchanged).
             _args: currently unused parameters.
             _kwargs: currently unused parameters.
         """
-        return convert_data_type(self, output_type=output_type, dtype=dtype, wrap_sequence=True)[0]
+        return convert_data_type(self, output_type=output_type, dtype=dtype, device=device, wrap_sequence=True)[0]
 
     def set_array(self, src, non_blocking=False, *_args, **_kwargs):
         """
@@ -428,7 +422,7 @@ def astype(self, dtype, device=None, *_args, **_kwargs):
             out_type = np.ndarray
         else:
             out_type = None
-        return convert_data_type(self, output_type=out_type, device=device, dtype=dtype, wrap_sequence=True)[0]
+        return self.get_array(output_type=out_type, dtype=dtype, device=device)
 
     @property
     def affine(self) -> torch.Tensor:
@@ -443,6 +437,8 @@ def affine(self, d: NdarrayTensor) -> None:
     @property
     def pixdim(self):
         """Get the spacing"""
+        if self.is_batch:
+            return [affine_to_spacing(a) for a in self.affine]
         return affine_to_spacing(self.affine)
 
     def new_empty(self, size, dtype=None, device=None, requires_grad=False):


Test patch:
diff --git a/tests/test_meta_tensor.py b/tests/test_meta_tensor.py
--- a/tests/test_meta_tensor.py
+++ b/tests/test_meta_tensor.py
@@ -365,6 +365,7 @@ def test_indexing(self):
         # `is_batch==False`, should have first metadata and subset of first image.
         d = data[0, 0]
         self.check(d, ims[0][0], ids=False)
+        self.assertEqual(d.applied_operations, ims[0][0].applied_operations)
 
         # `is_batch==True`, should have all metadata and subset of all images.
         d = data[:, 0]
@@ -389,6 +390,7 @@ def test_indexing(self):
         self.assertEqual(len(d), len(ims))
         for _d, _im in zip(d, ims):
             self.check(_d, _im, ids=False)
+            self.assertEqual(_d.applied_operations, _im.applied_operations)
 
         # `is_batch==True`, tuple split along non-batch dim. Should have all metadata.
         d = data.unbind(-1)