Enhance ExtractDataKeyFromMetaKeyd to work with MetaTensor

monai/apps/reconstruction/transforms/dictionary.py

@@ -20,6 +20,7 @@
 from monai.apps.reconstruction.transforms.array import EquispacedKspaceMask, RandomKspaceMask
 from monai.config import DtypeLike, KeysCollection
 from monai.config.type_definitions import NdarrayOrTensor
+from monai.data.meta_tensor import MetaTensor
 from monai.transforms import InvertibleTransform
 from monai.transforms.croppad.array import SpatialCrop
 from monai.transforms.intensity.array import NormalizeIntensity
@@ -33,15 +34,36 @@ class ExtractDataKeyFromMetaKeyd(MapTransform):
     Moves keys from meta to data. It is useful when a dataset of paired samples
     is loaded and certain keys should be moved from meta to data.
 
+    This transform supports two modes:
+
+    1. When ``meta_key`` references a metadata dictionary in the data (e.g., when
+       ``image_only=False`` was used with ``LoadImaged``), the requested keys are
+       extracted directly from that dictionary.
+
+    2. When ``meta_key`` references a ``MetaTensor`` in the data (e.g., when
+       ``image_only=True`` was used with ``LoadImaged``), the requested keys are
+       extracted from its ``.meta`` attribute.
+
     Args:
         keys: keys to be transferred from meta to data
-        meta_key: the meta key where all the meta-data is stored
+        meta_key: the key in the data dictionary where the metadata source is
+            stored. This can be either a metadata dictionary or a ``MetaTensor``.
         allow_missing_keys: don't raise exception if key is missing
 
     Example:
         When the fastMRI dataset is loaded, "kspace" is stored in the data dictionary,
         but the ground-truth image with the key "reconstruction_rss" is stored in the meta data.
         In this case, ExtractDataKeyFromMetaKeyd moves "reconstruction_rss" to data.
+
+        When ``LoadImaged`` is used with ``image_only=True`` (the default), the loaded
+        data is a ``MetaTensor`` with metadata accessible via ``.meta``. In this case,
+        set ``meta_key`` to the key of the ``MetaTensor`` itself::
+
+            li = LoadImaged(keys="image")  # image_only=True by default
+            dat = li({"image": "image.nii"})
+            e = ExtractDataKeyFromMetaKeyd("filename_or_obj", meta_key="image")
+            dat = e(dat)
+            assert dat["image"].meta["filename_or_obj"] == dat["filename_or_obj"]
     """
 
     def __init__(self, keys: KeysCollection, meta_key: str, allow_missing_keys: bool = False) -> None:
@@ -58,9 +80,18 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, T
             the new data dictionary
         """
         d = dict(data)
+        meta_obj = d[self.meta_key]
+
+        # If meta_key references a MetaTensor, extract from its .meta attribute;
+        # otherwise treat it as a metadata dictionary directly.
+        if isinstance(meta_obj, MetaTensor):
+            meta_dict = meta_obj.meta
+        else:
+            meta_dict = meta_obj
+
         for key in self.keys:
-            if key in d[self.meta_key]:
-                d[key] = d[self.meta_key][key]  # type: ignore
+            if key in meta_dict:
+                d[key] = meta_dict[key]  # type: ignore
             elif not self.allow_missing_keys:
                 raise KeyError(
                     f"Key `{key}` of transform `{self.__class__.__name__}` was missing in the meta data"

tests/apps/reconstruction/transforms/test_extract_data_key_from_meta_keyd.py

@@ -0,0 +1,116 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from monai.apps.reconstruction.transforms.dictionary import ExtractDataKeyFromMetaKeyd
+from monai.data import MetaTensor
+
+
+class TestExtractDataKeyFromMetaKeyd(unittest.TestCase):
+    """Tests for ExtractDataKeyFromMetaKeyd covering both dict-based and MetaTensor-based metadata."""
+
+    def test_extract_from_dict(self):
+        """Test extracting keys from a plain metadata dictionary (image_only=False scenario)."""
+        data = {
+            "image": torch.zeros(1, 2, 2),
+            "image_meta_dict": {"filename_or_obj": "image.nii", "spatial_shape": [2, 2]},
+        }
+        transform = ExtractDataKeyFromMetaKeyd(keys="filename_or_obj", meta_key="image_meta_dict")
+        result = transform(data)
+        self.assertIn("filename_or_obj", result)
+        self.assertEqual(result["filename_or_obj"], "image.nii")
+        self.assertEqual(result["image_meta_dict"]["filename_or_obj"], result["filename_or_obj"])
+
+    def test_extract_from_metatensor(self):
+        """Test extracting keys from a MetaTensor's .meta attribute (image_only=True scenario)."""
+        meta = {"filename_or_obj": "image.nii", "spatial_shape": [2, 2]}
+        mt = MetaTensor(torch.zeros(1, 2, 2), meta=meta)
+        data = {"image": mt}
+        transform = ExtractDataKeyFromMetaKeyd(keys="filename_or_obj", meta_key="image")
+        result = transform(data)
+        self.assertIn("filename_or_obj", result)
+        self.assertEqual(result["filename_or_obj"], "image.nii")
+        self.assertEqual(result["image"].meta["filename_or_obj"], result["filename_or_obj"])
+
+    def test_extract_multiple_keys_from_metatensor(self):
+        """Test extracting multiple keys from a MetaTensor."""
+        meta = {"filename_or_obj": "image.nii", "spatial_shape": [2, 2], "affine": "identity"}
+        mt = MetaTensor(torch.zeros(1, 2, 2), meta=meta)
+        data = {"image": mt}
+        transform = ExtractDataKeyFromMetaKeyd(keys=["filename_or_obj", "spatial_shape"], meta_key="image")
+        result = transform(data)
+        self.assertIn("filename_or_obj", result)
+        self.assertIn("spatial_shape", result)
+        self.assertEqual(result["filename_or_obj"], "image.nii")
+        self.assertEqual(result["spatial_shape"], [2, 2])
+
+    def test_extract_multiple_keys_from_dict(self):
+        """Test extracting multiple keys from a plain dictionary."""
+        data = {
+            "image": torch.zeros(1, 2, 2),
+            "image_meta_dict": {"filename_or_obj": "image.nii", "spatial_shape": [2, 2]},
+        }
+        transform = ExtractDataKeyFromMetaKeyd(keys=["filename_or_obj", "spatial_shape"], meta_key="image_meta_dict")
+        result = transform(data)
+        self.assertIn("filename_or_obj", result)
+        self.assertIn("spatial_shape", result)
+        self.assertEqual(result["filename_or_obj"], "image.nii")
+        self.assertEqual(result["spatial_shape"], [2, 2])
+
+    def test_missing_key_raises(self):
+        """Test that a missing key raises KeyError when allow_missing_keys=False."""
+        meta = {"filename_or_obj": "image.nii"}
+        mt = MetaTensor(torch.zeros(1, 2, 2), meta=meta)
+        data = {"image": mt}
+        transform = ExtractDataKeyFromMetaKeyd(keys="nonexistent_key", meta_key="image")
+        with self.assertRaises(KeyError):
+            transform(data)
+
+    def test_missing_key_allowed_metatensor(self):
+        """Test that a missing key is silently skipped when allow_missing_keys=True with MetaTensor."""
+        meta = {"filename_or_obj": "image.nii"}
+        mt = MetaTensor(torch.zeros(1, 2, 2), meta=meta)
+        data = {"image": mt}
+        transform = ExtractDataKeyFromMetaKeyd(keys="nonexistent_key", meta_key="image", allow_missing_keys=True)
+        result = transform(data)
+        self.assertNotIn("nonexistent_key", result)
+
+    def test_missing_key_allowed_dict(self):
+        """Test that a missing key is silently skipped when allow_missing_keys=True with dict."""
+        data = {
+            "image": torch.zeros(1, 2, 2),
+            "image_meta_dict": {"filename_or_obj": "image.nii"},
+        }
+        transform = ExtractDataKeyFromMetaKeyd(
+            keys="nonexistent_key", meta_key="image_meta_dict", allow_missing_keys=True
+        )
+        result = transform(data)
+        self.assertNotIn("nonexistent_key", result)
+
+    def test_original_data_preserved_metatensor(self):
+        """Test that the original MetaTensor remains in the data dictionary."""
+        meta = {"filename_or_obj": "image.nii"}
+        mt = MetaTensor(torch.ones(1, 2, 2), meta=meta)
+        data = {"image": mt}
+        transform = ExtractDataKeyFromMetaKeyd(keys="filename_or_obj", meta_key="image")
+        result = transform(data)
+        self.assertIn("image", result)
+        self.assertIsInstance(result["image"], MetaTensor)
+        self.assertTrue(torch.equal(result["image"], mt))
+
+
+if __name__ == "__main__":
+    unittest.main()