Enhance ExtractDataKeyFromMetaKeyd to work with MetaTensor#8772
Enhance ExtractDataKeyFromMetaKeyd to work with MetaTensor#8772haoyu-haoyu wants to merge 1 commit intoProject-MONAI:devfrom
Conversation
📝 WalkthroughWalkthroughExtractDataKeyFromMetaKeyd was updated to support MetaTensor metadata: it now imports MetaTensor, resolves the object at Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
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. Comment |
There was a problem hiding this comment.
🧹 Nitpick comments (4)
monai/apps/reconstruction/transforms/dictionary.py (2)
47-51: Missing Raises section in docstring.The
__call__method raisesKeyErrorwhen a requested key is absent andallow_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_keyreferences a value that is neitherMetaTensornordict, the code will fail at line 93 with an unclearTypeError. 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_keyitself is not present in the data dictionary. This would raiseKeyErroratd[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_raisesonly 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
📒 Files selected for processing (2)
monai/apps/reconstruction/transforms/dictionary.pytests/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>
9740a52 to
1b764f2
Compare
There was a problem hiding this comment.
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
📒 Files selected for processing (2)
monai/apps/reconstruction/transforms/dictionary.pytests/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
| 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)) | ||
|
|
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
find . -type f -name "*.py" | xargs rg "class ExtractDataKeyFromMetaKeyd" -A 30Repository: Project-MONAI/MONAI
Length of output: 3258
🏁 Script executed:
rg "class ExtractDataKeyFromMetaKeyd" -A 100 ./monai/apps/reconstruction/transforms/dictionary.py | head -120Repository: 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.
Fixes #7562
Description
Enhances
ExtractDataKeyFromMetaKeydto support extracting metadata fromMetaTensorobjects, in addition to plain metadata dictionaries.Before: Only worked with metadata dictionaries (
image_only=False):After: Also works with MetaTensor (
image_only=True, the default):Changes
monai/apps/reconstruction/transforms/dictionary.py:MetaTensorimportExtractDataKeyFromMetaKeyd.__call__()to detect ifmeta_keyreferences aMetaTensorand extract from its.metaattributetests/apps/reconstruction/transforms/test_extract_data_key_from_meta_keyd.py(new):allow_missing_keys), and data preservationTesting
Signed-off-by: haoyu-haoyu haoyu-haoyu@users.noreply.github.com