Skip to content

Enhance ExtractDataKeyFromMetaKeyd to work with MetaTensor#8772

Open
haoyu-haoyu wants to merge 2 commits intoProject-MONAI:devfrom
haoyu-haoyu:fix/enhance-extract-data-key-from-meta-keyd-metatensor
Open

Enhance ExtractDataKeyFromMetaKeyd to work with MetaTensor#8772
haoyu-haoyu wants to merge 2 commits intoProject-MONAI:devfrom
haoyu-haoyu:fix/enhance-extract-data-key-from-meta-keyd-metatensor

Conversation

@haoyu-haoyu
Copy link

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
Copy link
Contributor

coderabbitai bot commented Mar 11, 2026

📝 Walkthrough

Walkthrough

ExtractDataKeyFromMetaKeyd was updated to support MetaTensor metadata: it now imports MetaTensor, resolves the object at meta_key, and if that object is a MetaTensor uses its .meta dict (otherwise uses the dict as-is) when extracting keys to copy into data. The call logic and KeyError behavior (when allow_missing_keys=False) were adjusted to reference the resolved metadata dict. Docstring expanded with examples for both dictionary- and MetaTensor-based modes. A new test module covers dict and MetaTensor extraction, multi-key cases, missing-key handling, 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 summarizes the main enhancement: adding MetaTensor support to ExtractDataKeyFromMetaKeyd.
Description check ✅ Passed Description is comprehensive, following the template with before/after examples, clear change documentation, testing notes, and backward compatibility confirmation.
Linked Issues check ✅ Passed Changes fully address #7562: added MetaTensor detection, extraction from .meta attribute, and comprehensive test coverage for both dict and MetaTensor modes.
Out of Scope Changes check ✅ Passed All changes are directly scoped to #7562: modified ExtractDataKeyFromMetaKeyd implementation and added unit tests; no unrelated alterations present.
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
📝 Coding Plan
  • Generate coding plan for human review comments

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

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

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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

When LoadImaged is used with image_only=True (the default), the loaded
data is a MetaTensor with metadata accessible via .meta attribute.
Previously, ExtractDataKeyFromMetaKeyd could only extract keys from
metadata dictionaries (image_only=False scenario).

This change adds support for MetaTensor by detecting if the meta_key
references a MetaTensor instance and extracting from its .meta attribute
instead of treating it as a plain dictionary.

Fixes Project-MONAI#7562

Signed-off-by: haoyu-haoyu <haoyu-haoyu@users.noreply.github.com>
Signed-off-by: SexyERIC0723 <haoyuwang144@gmail.com>
@haoyu-haoyu haoyu-haoyu force-pushed the fix/enhance-extract-data-key-from-meta-keyd-metatensor branch from 9740a52 to 1b764f2 Compare March 14, 2026 02:02
Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
tests/apps/reconstruction/transforms/test_extract_data_key_from_meta_keyd.py (1)

23-105: Docstrings are not in required Google-style sections.

Definitions have docstrings, but they don’t document Args/Returns/Raises in Google-style format.

As per coding guidelines, "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."

🤖 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 23 - 105, The test methods (e.g., test_extract_from_dict,
test_extract_from_metatensor, test_extract_multiple_keys_from_metatensor,
test_missing_key_raises, test_missing_key_allowed_metatensor,
test_missing_key_allowed_dict, test_original_data_preserved_metatensor)
currently use free-form docstrings; update each test docstring to Google-style
sections: add Args: (describe inputs like data, transform), Returns: (if any)
and Raises: (e.g., KeyError for test_missing_key_raises) where appropriate,
following the project's docstring guidelines, ensuring each docstring documents
arguments, expected return/side-effect, and exceptions raised for clarity and
consistency with ExtractDataKeyFromMetaKeyd usage.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In
`@tests/apps/reconstruction/transforms/test_extract_data_key_from_meta_keyd.py`:
- Around line 103-113: In test_original_data_preserved_metatensor change the
final assertion to check object identity instead of tensor equality: replace the
torch.equal-based assertion with an assertIs on result["image"] and mt so the
test verifies the transform (ExtractDataKeyFromMetaKeyd) preserved the original
MetaTensor object reference rather than just equal values.

---

Nitpick comments:
In
`@tests/apps/reconstruction/transforms/test_extract_data_key_from_meta_keyd.py`:
- Around line 23-105: The test methods (e.g., test_extract_from_dict,
test_extract_from_metatensor, test_extract_multiple_keys_from_metatensor,
test_missing_key_raises, test_missing_key_allowed_metatensor,
test_missing_key_allowed_dict, test_original_data_preserved_metatensor)
currently use free-form docstrings; update each test docstring to Google-style
sections: add Args: (describe inputs like data, transform), Returns: (if any)
and Raises: (e.g., KeyError for test_missing_key_raises) where appropriate,
following the project's docstring guidelines, ensuring each docstring documents
arguments, expected return/side-effect, and exceptions raised for clarity and
consistency with ExtractDataKeyFromMetaKeyd usage.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: d27d7bff-9213-45b1-ad0a-df29966a5710

📥 Commits

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

📒 Files selected for processing (2)
  • monai/apps/reconstruction/transforms/dictionary.py
  • tests/apps/reconstruction/transforms/test_extract_data_key_from_meta_keyd.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • monai/apps/reconstruction/transforms/dictionary.py

Comment on lines +103 to +113
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))

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

find . -type f -name "*.py" | xargs rg "class ExtractDataKeyFromMetaKeyd" -A 30

Repository: Project-MONAI/MONAI

Length of output: 3258


🏁 Script executed:

rg "class ExtractDataKeyFromMetaKeyd" -A 100 ./monai/apps/reconstruction/transforms/dictionary.py | head -120

Repository: Project-MONAI/MONAI

Length of output: 4638


Line 112 does not verify "original object preserved."

torch.equal(...) checks tensor value equality, not object identity. Since the transform preserves the original MetaTensor object reference (via shallow dict copy), use assertIs() to verify:

Proposed test fix
-        self.assertTrue(torch.equal(result["image"], mt))
+        self.assertIs(result["image"], mt)
🤖 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 103 - 113, In test_original_data_preserved_metatensor change the
final assertion to check object identity instead of tensor equality: replace the
torch.equal-based assertion with an assertIs on result["image"] and mt so the
test verifies the transform (ExtractDataKeyFromMetaKeyd) preserved the original
MetaTensor object reference rather than just equal values.

Signed-off-by: SexyERIC0723 <haoyuwang144@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Enhance ExtractDataKeyFromMetaKeyd to work with MetaTensor

1 participant