Skip to content

Enhance/metatensor extraction #8754

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
skdas20 wants to merge 3 commits into
base: dev
Choose a base branch
from
Open

Enhance/metatensor extraction #8754

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
85ed7a9
fix: support more dtypes in pad_nd function
skdas20 Feb 22, 2026
8fdad6d
improve: add helpful error message for missing tensorboard package
skdas20 Feb 22, 2026
5f84f35
enhance: support MetaTensor in ExtractDataKeyFromMetaKeyd
skdas20 Feb 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions monai/apps/reconstruction/transforms/dictionary.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,16 +37,19 @@ class ExtractDataKeyFromMetaKeyd(MapTransform):
keys: keys to be transferred from meta to data
meta_key: the meta key where all the meta-data is stored
allow_missing_keys: don't raise exception if key is missing
image_only: if True, only extract metadata from MetaTensor images to avoid duplication

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.
For MetaTensor objects, setting image_only=True prevents extracting redundant metadata.
Comment on lines +40 to +46
Copy link
Contributor

@coderabbitai coderabbitai bot Feb 22, 2026

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟡 Minor

Docstring for image_only inverts the actual behavior.

Line 40 says "only extract metadata from MetaTensor images", but the real effect is the opposite: MetaTensor values are skipped (not extracted). The example note on line 46 is also vague. Consider:

-        image_only: if True, only extract metadata from MetaTensor images to avoid duplication
+        image_only: if True, skip keys whose value is a MetaTensor (leaves them in meta_key only).
-        For MetaTensor objects, setting image_only=True prevents extracting redundant metadata.
+        When image_only=True, MetaTensor-valued keys are not moved to the top-level data dict.
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
image_only: if True, only extract metadata from MetaTensor images to avoid duplication
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.
For MetaTensor objects, setting image_only=True prevents extracting redundant metadata.
image_only: if True, skip keys whose value is a MetaTensor (leaves them in meta_key only).
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 image_only=True, MetaTensor-valued keys are not moved to the top-level data dict.
🤖 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 40 - 46, The
docstring for ExtractDataKeyFromMetaKeyd incorrectly describes image_only;
update the parameter description and example to reflect the actual behavior:
when image_only=True the transform skips/excludes MetaTensor values (does not
extract their metadata) and only extracts metadata from non-MetaTensor entries,
and when image_only=False it will extract from MetaTensor items as well; edit
the parameter line mentioning image_only and the example text to explicitly
state this skip behavior and show a concrete fastMRI example where
"reconstruction_rss" remains in meta when image_only=True and is moved to data
when image_only=False.

"""

def __init__(self, keys: KeysCollection, meta_key: str, allow_missing_keys: bool = False) -> None:
def __init__(self, keys: KeysCollection, meta_key: str, allow_missing_keys: bool = False, image_only: bool = False) -> None:
MapTransform.__init__(self, keys, allow_missing_keys)
self.meta_key = meta_key
self.image_only = image_only
Comment on lines +49 to +52
Copy link
Contributor

@coderabbitai coderabbitai bot Feb 22, 2026

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

# Search for test files and test coverage related to ExtractDataKeyFromMetaKeyd
find . -type f -name "*test*dictionary*" | head -20

Repository: Project-MONAI/MONAI

Length of output: 45

🏁 Script executed:

# Search for references to ExtractDataKeyFromMetaKeyd in the codebase
rg "ExtractDataKeyFromMetaKeyd" --type py -l

Repository: Project-MONAI/MONAI

Length of output: 113

🏁 Script executed:

# Search for image_only parameter usage/tests
rg "image_only" --type py -B 2 -A 2

Repository: Project-MONAI/MONAI

Length of output: 50375

🏁 Script executed:

# Examine the implementation file to understand the parameter
head -100 monai/apps/reconstruction/transforms/dictionary.py

Repository: Project-MONAI/MONAI

Length of output: 4812

🏁 Script executed:

# Look for test files in reconstruction directory
find . -path "*reconstruction*test*" -type f -name "*.py"

Repository: Project-MONAI/MONAI

Length of output: 485

🏁 Script executed:

# Search for any references to ExtractDataKeyFromMetaKeyd in the entire codebase
rg "ExtractDataKeyFromMetaKeyd" -B 3 -A 3

Repository: Project-MONAI/MONAI

Length of output: 2000

🏁 Script executed:

# Examine the specific transform file more carefully to see the full implementation
sed -n '40,100p' monai/apps/reconstruction/transforms/dictionary.py

Repository: Project-MONAI/MONAI

Length of output: 3131

🏁 Script executed:

# Search for any tests that reference ExtractDataKeyFromMetaKeyd
rg "ExtractDataKeyFromMetaKeyd" tests/

Repository: Project-MONAI/MONAI

Length of output: 45

🏁 Script executed:

# Check if MetaTensor is imported in the dictionary.py file
head -30 monai/apps/reconstruction/transforms/dictionary.py

Repository: Project-MONAI/MONAI

Length of output: 1342

🏁 Script executed:

# Check the full imports and look for MetaTensor import
cat monai/apps/reconstruction/transforms/dictionary.py | head -50

Repository: Project-MONAI/MONAI

Length of output: 2417

🏁 Script executed:

# Search for MetaTensor import in the file
rg "MetaTensor" monai/apps/reconstruction/transforms/dictionary.py

Repository: Project-MONAI/MONAI

Length of output: 420

🏁 Script executed:

# Check if MetaTensor needs to be imported - search where it's defined
rg "class MetaTensor" monai/

Repository: Project-MONAI/MONAI

Length of output: 129

🏁 Script executed:

# Look at the full __call__ method to understand the context better
sed -n '50,80p' monai/apps/reconstruction/transforms/dictionary.py

Repository: Project-MONAI/MONAI

Length of output: 1288

🏁 Script executed:

# Check the entire dictionary.py file for MetaTensor imports
cat monai/apps/reconstruction/transforms/dictionary.py | grep -n "MetaTensor\|import"

Repository: Project-MONAI/MONAI

Length of output: 1175

🏁 Script executed:

# Check if MetaTensor is used anywhere in monai.data that we can import from
rg "from monai.data" monai/apps/reconstruction/transforms/dictionary.py

Repository: Project-MONAI/MONAI

Length of output: 45

🏁 Script executed:

# Look at similar dictionary transforms to see how they handle MetaTensor
rg "MetaTensor" monai/transforms/io/dictionary.py -B 2 -A 2 | head -20

Repository: Project-MONAI/MONAI

Length of output: 45

🏁 Script executed:

# Check what MetaTensor import is needed
rg "from.*MetaTensor" monai/transforms/ | head -5

Repository: Project-MONAI/MONAI

Length of output: 605

🏁 Script executed:

# Verify that MetaTensor is referenced but not imported by checking if code would fail
python3 << 'EOF'
import ast
import sys

with open('monai/apps/reconstruction/transforms/dictionary.py', 'r') as f:
    content = f.read()
    
# Simple check: look for MetaTensor in code but not in imports
has_metatensor_import = 'from monai.data.meta_tensor import MetaTensor' in content
has_metatensor_usage = 'isinstance(extracted_value, MetaTensor)' in content

print(f"MetaTensor imported: {has_metatensor_import}")
print(f"MetaTensor used in code: {has_metatensor_usage}")
print(f"Will cause NameError if image_only=True is used: {has_metatensor_usage and not has_metatensor_import}")
EOF

Repository: Project-MONAI/MONAI

Length of output: 173

🏁 Script executed:

# Check if there are ANY tests for ExtractDataKeyFromMetaKeyd
find tests/ -type f -name "*.py" -exec grep -l "ExtractDataKeyFromMetaKeyd" {} \;

Repository: Project-MONAI/MONAI

Length of output: 45

Add missing MetaTensor import and implement unit tests for the image_only parameter.

MetaTensor is referenced in line 69 but not imported—the code will raise NameError when image_only=True. Import from monai.data.meta_tensor. Additionally, add test coverage for both image_only=False (default) and image_only=True with MetaTensor and plain tensor values to verify the skip and extract paths.

🤖 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 49 - 52,
Import MetaTensor from monai.data.meta_tensor and update the transform to
reference it properly so using image_only=True doesn't raise NameError: add
"from monai.data.meta_tensor import MetaTensor" near other imports and ensure
the __init__ signature (def __init__(..., image_only: bool = False)) and logic
that checks MetaTensor use the imported name. Add unit tests for the transform
covering both image_only=False and image_only=True: create test cases using a
MetaTensor value and a plain torch tensor to verify the extraction path when
image_only=False and the skip/retain behavior when image_only=True; assert the
meta_key is extracted for MetaTensor and unchanged for plain tensors. Ensure
tests exercise the transform class/method that performs meta extraction (the
class containing __init__ and the method that handles meta_key extraction) and
include both positive and negative assertions for MetaTensor vs plain tensor
inputs.


def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, Tensor]:
"""
Expand All @@ -60,7 +63,12 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, T
d = dict(data)
for key in self.keys:
if key in d[self.meta_key]:
d[key] = d[self.meta_key][key] # type: ignore
extracted_value = d[self.meta_key][key]
# When image_only is True, skip if the extracted value is a MetaTensor
# to preserve metadata associations
if self.image_only and isinstance(extracted_value, MetaTensor):
continue
d[key] = extracted_value # type: ignore
Comment on lines +66 to +71
Copy link
Contributor

@coderabbitai coderabbitai bot Feb 22, 2026

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🔴 Critical

MetaTensor is not imported — NameError at runtime.

MetaTensor is referenced at line 69 but never imported in this file. When image_only=True and a MetaTensor value is encountered, the code raises NameError: name 'MetaTensor' is not defined. This is confirmed by Ruff F821.

🐛 Proposed fix

Add the import after the existing monai.data / monai.transforms imports:

+from monai.data.meta_tensor import MetaTensor
🧰 Tools
🪛 Ruff (0.15.1)

[error] 69-69: Undefined name MetaTensor

(F821)

🤖 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 66 - 71, The
code references MetaTensor in the dictionary extraction logic (used in the block
that checks self.image_only and calls isinstance(extracted_value, MetaTensor)
when populating d[key]) but MetaTensor is not imported; fix by adding the proper
import (e.g., add "from monai.data import MetaTensor" alongside the existing
monai imports) so that the isinstance check in that extraction loop (the code
around extracted_value = d[self.meta_key][key] and the self.image_only branch)
does not raise NameError at runtime.

Comment on lines +67 to +71
Copy link
Contributor

@coderabbitai coderabbitai bot Feb 22, 2026

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major

Silent continue silently drops expected keys when image_only=True.

If image_only=True and a requested key's value is a MetaTensor, the key is never written into d. Any downstream transform accessing d[key] will get a KeyError with no warning. Either document this contract explicitly (the key stays in d[self.meta_key] only) or raise/log when a key is silently skipped.

Additionally, the parameter name image_only is semantically misleading — MetaTensor is the image object, so image_only=True skipping MetaTensors reads backwards. A name like skip_meta_tensor or non_tensor_only would be less ambiguous.

🧰 Tools
🪛 Ruff (0.15.1)

[error] 69-69: Undefined name MetaTensor

(F821)

🤖 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 67 - 71, The
current loop in the transform silently skips writing d[key] when self.image_only
is True and extracted_value is a MetaTensor (using isinstance(extracted_value,
MetaTensor)), which can cause downstream KeyError; change this behavior in the
method that builds d (referencing self.image_only, extracted_value, MetaTensor,
d[key], and self.meta_key) to either (1) explicitly log a warning (processLogger
or logger.warning) indicating the key was skipped and why, or (2) raise a clear
exception (e.g., ValueError) when image_only=True would drop a MetaTensor so
callers notice, and update the docstring to explain the contract; additionally,
rename the parameter self.image_only to a clearer name such as
self.skip_meta_tensor (and update all usages and docstrings) to avoid the
semantic confusion that MetaTensor is the image object.

elif not self.allow_missing_keys:
raise KeyError(
f"Key `{key}` of transform `{self.__class__.__name__}` was missing in the meta data"
Expand Down
5 changes: 5 additions & 0 deletions monai/handlers/tensorboard_handlers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,11 @@ class TensorBoardHandler:

def __init__(self, summary_writer: SummaryWriter | SummaryWriterX | None = None, log_dir: str = "./runs"):
if summary_writer is None:
if SummaryWriter is None:
raise RuntimeError(
"TensorBoardHandler requires tensorboard to be installed. "
"Please install it with: pip install tensorboard"
)
Comment on lines +53 to +57
Copy link
Contributor

@coderabbitai coderabbitai bot Feb 22, 2026

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🔴 Critical

SummaryWriter is None is always False — the guard is dead code.

optional_import returns the imported module and a boolean flag indicating whether the import is successful. When tensorboard is missing, importing-related exceptions are stored and raised lazily when attempting to use the failed-to-import module — the returned object is a stub, never None. So SummaryWriter is None never triggers; the RuntimeError is never raised and the helpful message never shown. The stub's OptionalImportError fires instead at line 58.

Fix: retain the success flag from optional_import and check it.

🐛 Proposed fix

At line 34, capture the boolean:

-    SummaryWriter, _ = optional_import("torch.utils.tensorboard", name="SummaryWriter")
+    SummaryWriter, _has_summary_writer = optional_import("torch.utils.tensorboard", name="SummaryWriter")

Then in __init__:

-            if SummaryWriter is None:
+            if not _has_summary_writer:
                 raise RuntimeError(
                     "TensorBoardHandler requires tensorboard to be installed. "
                     "Please install it with: pip install tensorboard"
                 )
🧰 Tools
🪛 Ruff (0.15.1)

[warning] 54-57: Avoid specifying long messages outside the exception class

(TRY003)

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@monai/handlers/tensorboard_handlers.py` around lines 53 - 57, Capture the
boolean success flag returned by optional_import when importing SummaryWriter
(e.g., SummaryWriter, summary_writer_available = optional_import(...)) and store
that alongside the SummaryWriter symbol; then in TensorBoardHandler.__init__
check the boolean flag (summary_writer_available) instead of testing
"SummaryWriter is None" and raise the RuntimeError with the helpful install
message if the flag is False so the guard triggers before the
OptionalImportError stub is invoked.

self._writer = SummaryWriter(log_dir=log_dir)
self.internal_writer = True
else:
Expand Down
7 changes: 1 addition & 6 deletions monai/transforms/croppad/functional.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,12 +96,7 @@ def pad_nd(
return _np_pad(img, pad_width=to_pad, mode=mode, **kwargs)
try:
_pad = _np_pad
if mode in {"constant", "reflect", "edge", "replicate", "wrap", "circular"} and img.dtype not in {
torch.int16,
torch.int64,
torch.bool,
torch.uint8,
}:
if mode in {"constant", "reflect", "edge", "replicate", "wrap", "circular"}:
_pad = _pt_pad
return _pad(img, pad_width=to_pad, mode=mode, **kwargs)
except (ValueError, TypeError, RuntimeError) as err:
Expand Down
Loading