Enhance ExtractDataKeyFromMetaKeyd to work with MetaTensor

[haoyu-haoyu]

Fixes #7562

Description

Enhances ExtractDataKeyFromMetaKeyd to support extracting metadata from MetaTensor objects, in addition to plain metadata dictionaries.

Before: Only worked with metadata dictionaries (image_only=False):

li = LoadImaged(image_only=False)
dat = li({"image": "image.nii"})  # {"image": tensor, "image_meta_dict": dict}
e = ExtractDataKeyFromMetaKeyd("filename_or_obj", meta_key="image_meta_dict")
dat = e(dat)  # extracts from dict

After: Also works with MetaTensor (image_only=True, the default):

li = LoadImaged()  # image_only=True by default
dat = li({"image": "image.nii"})  # {"image": MetaTensor}
e = ExtractDataKeyFromMetaKeyd("filename_or_obj", meta_key="image")
dat = e(dat)  # extracts from MetaTensor.meta
assert dat["image"].meta["filename_or_obj"] == dat["filename_or_obj"]

Changes

1. monai/apps/reconstruction/transforms/dictionary.py:

   • Added MetaTensor import
   • Modified ExtractDataKeyFromMetaKeyd.__call__() to detect if meta_key references a MetaTensor and extract from its .meta attribute
   • Updated docstring with both usage modes and examples

2. tests/apps/reconstruction/transforms/test_extract_data_key_from_meta_keyd.py (new):

   • 8 test cases covering: dict extraction, MetaTensor extraction, multiple keys, missing keys (with/without allow_missing_keys), and data preservation

Testing

• New unit tests for both dict-based and MetaTensor-based extraction
• Tests for edge cases (missing keys, allow_missing_keys)
• Backward compatible — existing dict-based usage unchanged

Signed-off-by: haoyu-haoyu haoyu-haoyu@users.noreply.github.com

[coderabbitai]

📝 Walkthrough

Walkthrough

ExtractDataKeyFromMetaKeyd now supports extracting metadata from MetaTensor objects. Runtime detection checks if the metadata source is a MetaTensor and accesses its .meta attribute; otherwise it uses the provided dictionary. The docstring documents both modes with examples. A new test module validates extraction from dict and MetaTensor metadata, multi-key extraction, missing-key handling with and without allow_missing_keys, and MetaTensor preservation.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title clearly describes the main enhancement: adding MetaTensor support to ExtractDataKeyFromMetaKeyd.
Description check	✅ Passed	PR description comprehensively covers the changes, includes before/after examples, details modifications, and documents testing coverage.
Linked Issues check	✅ Passed	Changes directly address #7562: ExtractDataKeyFromMetaKeyd now detects and extracts from MetaTensor.meta when referenced via meta_key, enabling use with image_only=True.
Out of Scope Changes check	✅ Passed	All changes are scoped to the stated objective: enhancing ExtractDataKeyFromMetaKeyd with MetaTensor support and adding comprehensive test coverage.
Docstring Coverage	✅ Passed	Docstring coverage is 90.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (4)

monai/apps/reconstruction/transforms/dictionary.py (2)

47-51: Missing Raises section in docstring.

The __call__ method raises KeyError when a requested key is absent and allow_missing_keys=False. Per coding guidelines, docstrings should document raised exceptions.

📝 Proposed docstring addition

         allow_missing_keys: don't raise exception if key is missing
+
+    Raises:
+        KeyError: If ``meta_key`` is not found in the data dictionary, or if a
+            requested key is missing from the metadata and ``allow_missing_keys``
+            is False.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@monai/apps/reconstruction/transforms/dictionary.py` around lines 47 - 51, The
docstring for the transform is missing a Raises section documenting that
__call__ can raise KeyError when a requested key is absent and
allow_missing_keys is False; update the docstring for the class or the __call__
method (reference symbols: __call__, keys, meta_key, allow_missing_keys) to add
a "Raises" section that clearly states KeyError is raised in that situation and
any conditions under which it is not raised.

---

83-90: Consider validating meta_obj type.

If meta_key references a value that is neither MetaTensor nor dict, the code will fail at line 93 with an unclear TypeError. Adding a type check would provide a clearer error message.

🛡️ Proposed defensive check

         if isinstance(meta_obj, MetaTensor):
             meta_dict = meta_obj.meta
-        else:
+        elif isinstance(meta_obj, dict):
             meta_dict = meta_obj
+        else:
+            raise TypeError(
+                f"meta_key `{self.meta_key}` must reference a MetaTensor or dict, "
+                f"got {type(meta_obj).__name__}"
+            )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@monai/apps/reconstruction/transforms/dictionary.py` around lines 83 - 90,
meta_obj retrieved by meta_key may be neither a MetaTensor nor a dict, so add an
explicit type check after obtaining meta_obj (in the block that currently tests
isinstance(meta_obj, MetaTensor)) to validate allowed types: if meta_obj is a
MetaTensor set meta_dict = meta_obj.meta; elif it's a dict set meta_dict =
meta_obj; else raise a clear TypeError referencing meta_key and the actual type
of meta_obj. Update the code paths around meta_key / meta_obj / meta_dict to use
this validated meta_dict afterwards.

tests/apps/reconstruction/transforms/test_extract_data_key_from_meta_keyd.py (2)

119-122: Consider adding test for missing meta_key.

No test covers the case when meta_key itself is not present in the data dictionary. This would raise KeyError at d[self.meta_key].

🧪 Proposed edge case test

def test_missing_meta_key_raises(self):
    """Test that missing meta_key raises KeyError."""
    data = {"image": torch.zeros(1, 2, 2)}
    transform = ExtractDataKeyFromMetaKeyd(keys="filename_or_obj", meta_key="nonexistent_meta")
    with self.assertRaises(KeyError):
        transform(data)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/apps/reconstruction/transforms/test_extract_data_key_from_meta_keyd.py`
around lines 119 - 122, Add a test covering the case where the meta_key is
absent: create a test method (e.g., test_missing_meta_key_raises) that
constructs input data without the specified meta_key, instantiates
ExtractDataKeyFromMetaKeyd(keys="filename_or_obj", meta_key="nonexistent_meta"),
and asserts that calling transform(data) raises a KeyError; this validates
behavior when d[self.meta_key] is missing.

---

77-84: Missing test for KeyError with dict source.

test_missing_key_raises only tests with MetaTensor. Add a similar test for dict-based metadata to ensure symmetry.

🧪 Proposed additional test

def test_missing_key_raises_dict(self):
    """Test that a missing key raises KeyError when allow_missing_keys=False 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")
    with self.assertRaises(KeyError):
        transform(data)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/apps/reconstruction/transforms/test_extract_data_key_from_meta_keyd.py`
around lines 77 - 84, Add a parallel unit test to cover the dict-based metadata
case: create a test function (e.g., test_missing_key_raises_dict) that
constructs a plain tensor data dict with an accompanying meta dict key (e.g.,
"image_meta_dict") missing the requested meta field, instantiate
ExtractDataKeyFromMetaKeyd(keys="nonexistent_key", meta_key="image_meta_dict"),
and assert it raises KeyError when called; mirror the existing
test_missing_key_raises structure but use a raw torch tensor for "image" and a
dict for the meta source to ensure symmetry between MetaTensor and dict
handling.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@monai/apps/reconstruction/transforms/dictionary.py`:
- Around line 47-51: The docstring for the transform is missing a Raises section
documenting that __call__ can raise KeyError when a requested key is absent and
allow_missing_keys is False; update the docstring for the class or the __call__
method (reference symbols: __call__, keys, meta_key, allow_missing_keys) to add
a "Raises" section that clearly states KeyError is raised in that situation and
any conditions under which it is not raised.
- Around line 83-90: meta_obj retrieved by meta_key may be neither a MetaTensor
nor a dict, so add an explicit type check after obtaining meta_obj (in the block
that currently tests isinstance(meta_obj, MetaTensor)) to validate allowed
types: if meta_obj is a MetaTensor set meta_dict = meta_obj.meta; elif it's a
dict set meta_dict = meta_obj; else raise a clear TypeError referencing meta_key
and the actual type of meta_obj. Update the code paths around meta_key /
meta_obj / meta_dict to use this validated meta_dict afterwards.

In
`@tests/apps/reconstruction/transforms/test_extract_data_key_from_meta_keyd.py`:
- Around line 119-122: Add a test covering the case where the meta_key is
absent: create a test method (e.g., test_missing_meta_key_raises) that
constructs input data without the specified meta_key, instantiates
ExtractDataKeyFromMetaKeyd(keys="filename_or_obj", meta_key="nonexistent_meta"),
and asserts that calling transform(data) raises a KeyError; this validates
behavior when d[self.meta_key] is missing.
- Around line 77-84: Add a parallel unit test to cover the dict-based metadata
case: create a test function (e.g., test_missing_key_raises_dict) that
constructs a plain tensor data dict with an accompanying meta dict key (e.g.,
"image_meta_dict") missing the requested meta field, instantiate
ExtractDataKeyFromMetaKeyd(keys="nonexistent_key", meta_key="image_meta_dict"),
and assert it raises KeyError when called; mirror the existing
test_missing_key_raises structure but use a raw torch tensor for "image" and a
dict for the meta source to ensure symmetry between MetaTensor and dict
handling.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 8c508da7-ee91-44a3-88b1-2a1c5c11d542

📥 Commits

Reviewing files that changed from the base of the PR and between daaedaa and 9740a52.

📒 Files selected for processing (2)

• monai/apps/reconstruction/transforms/dictionary.py
• tests/apps/reconstruction/transforms/test_extract_data_key_from_meta_keyd.py