Add nnunet - #88
Draft
innat wants to merge 4 commits into
Draft
Add nnunet #88innat wants to merge 4 commits into
nnunet #88innat wants to merge 4 commits into
Conversation
innat
added
enhancement
📝 WalkthroughWalkthroughThis pull request introduces a comprehensive Keras-based nnU-Net implementation for medical image segmentation, spanning model architecture, dataset management, training orchestration, planning utilities, and preprocessing pipelines. The implementation includes reusable neural network building blocks, a configurable UNet model, dataset manifest/fingerprinting/preprocessing systems, cross-validation utilities, a training orchestrator with multi-output supervision support, and end-to-end I/O and normalization helpers. Changes
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive Keras 3 implementation of the nnU-Net workflow, including model architecture, dataset fingerprinting, and training orchestration. While the architecture is well-structured, there are critical memory management issues in the fingerprinting and normalization modules where accumulating entire datasets in memory will lead to Out-Of-Memory errors. Additionally, the augmentation pipeline contains inefficient layer instantiations that may cause memory leaks, and the U-Net blocks have incorrect serialization logic in their get_config methods. Finally, the preprocessing stage currently lacks multi-processing support, which will significantly bottleneck performance.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (16)
medicai/trainer/nnunet/pipeline.py (1)
56-68: Consider logging GPU detection failures for debugging.The silent fallback is acceptable for robustness, but logging a debug/warning message would help troubleshoot cases where auto-detection unexpectedly fails.
Proposed improvement
+import logging + +logger = logging.getLogger(__name__) + def _auto_detect_gpu_memory(default=8.0): ... try: result = subprocess.run( ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"], stdout=subprocess.PIPE, text=True, check=True, ) memories = [float(x) / 1024.0 for x in result.stdout.strip().split("\n") if x.strip()] if memories: return min(memories) - except Exception: - pass + except Exception as e: + logger.debug("GPU memory auto-detection failed, using default %.1fGB: %s", default, e) return float(default)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/pipeline.py` around lines 56 - 68, The GPU auto-detection block silently swallows exceptions; modify the except block to log the failure (e.g., using a module logger obtained via logging.getLogger(__name__)) so failures in subprocess.run/nvidia-smi parsing are visible for debugging. Specifically, in the try/except around subprocess.run/result/memories, catch Exception as e and call logger.debug or logger.warning with a clear message like "Failed to detect GPU memory via nvidia-smi" including the exception details and optionally result.stdout, then fall back to returning float(default) as before.medicai/trainer/nnunet/data/resampling.py (1)
33-36: Addstrict=Truetozip()for defensive validation.The assertion validates length equality, but adding
strict=Truetozip()provides an extra safeguard and satisfies the Ruff B905 warning.Proposed fix
assert len(original_spacing) == len( target_spacing ), f"Spacing length mismatch: {len(original_spacing)} vs {len(target_spacing)}" - return tuple(float(o) / float(t) for o, t in zip(original_spacing, target_spacing)) + return tuple(float(o) / float(t) for o, t in zip(original_spacing, target_spacing, strict=True))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/data/resampling.py` around lines 33 - 36, The length assertion is good but we should make the resampling ratio computation defensive by using zip(..., strict=True) to satisfy Ruff B905; update the return in resampling.py (the tuple comprehension that currently uses zip(original_spacing, target_spacing)) to call zip(original_spacing, target_spacing, strict=True) so any length mismatch raises immediately in zip as well (keep the existing assert in place).medicai/trainer/nnunet/README.md (1)
68-75: Minor style improvement for readability.Consider varying sentence structure to improve flow, though this is optional.
Suggested rewording
### Spacing Rules -- If `spacing` is present on an item, it is used. -- If a NIfTI or DICOM file omits manifest spacing, MedicAI tries to read spacing from file metadata. -- If spacing still cannot be determined, MedicAI falls back to isotropic spacing: +- When `spacing` is present on an item, that value takes precedence. +- For NIfTI or DICOM files without manifest spacing, MedicAI reads spacing from file metadata. +- As a final fallback when spacing cannot be determined, MedicAI uses isotropic spacing:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/README.md` around lines 68 - 75, The "Spacing Rules" section is slightly repetitive and can be reworded for smoother flow; update the README paragraph(s) under the "Spacing Rules" heading to vary sentence structure and improve readability by combining related points (e.g., state that an item’s explicit `spacing` is used first, then fall back to file metadata for NIfTI/DICOM, and finally to isotropic defaults), and present the default isotropic values clearly (2D: [1.0, 1.0], 3D: [1.0, 1.0, 1.0]) while keeping references to `spacing`, NIfTI, DICOM, and MedicAI intact.medicai/models/nnunet/unet.py (1)
258-263: Unused loop variablei.Proposed fix
- for i in range(n_pooling + 1): + for _ in range(n_pooling + 1): schedule.append(min(f, max_filters)) f *= 2🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/models/nnunet/unet.py` around lines 258 - 263, The for-loop that builds the filter schedule uses an unused loop variable `i`; update the loop to use a throwaway variable (e.g., `for _ in range(n_pooling + 1):`) so linters won't flag an unused variable while preserving the logic that appends min(f, max_filters) and doubles `f`; adjust the loop in the schedule construction that references `base_filters`, `n_pooling`, `max_filters`, `schedule`, and `f` accordingly.medicai/trainer/nnunet/data/dataset_fingerprint.py (2)
105-127: Consider addingstrict=Trueto zip to catch length mismatches.When
target_class_idsis provided, it may have a different length thanlabel_paths. Thezipwill silently truncate to the shorter length, potentially causing incorrect class statistics.Proposed fix
- for class_id, label_path in zip(class_ids, label_paths): + for class_id, label_path in zip(class_ids, label_paths, strict=True):Alternatively, if truncation is intentional, add a comment explaining the behavior.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/data/dataset_fingerprint.py` around lines 105 - 127, The loop using zip(class_ids, label_paths) in dataset_fingerprint.py can silently truncate when target_class_ids length differs from label_paths; modify DatasetFingerprint logic to validate lengths before iterating (e.g., check len(class_ids) == len(label_paths) when target_class_ids is provided and raise a clear error) or explicitly use strict=True in the zip to force a ValueError on mismatch; update the code around class_ids, label_paths and the for ... in zip(...) loop in the function where counts and total_voxels are computed (references: class_ids, label_paths, counts, total_voxels) and add a short comment explaining the chosen behavior.
291-296: Consider using unpacking syntax for list concatenation (optional).This is a minor style improvement suggested by the linter.
Proposed fix
if num_dims == 2: - median_spacing = [1.0] + median_spacing - median_size = [1] + median_size + median_spacing = [1.0, *median_spacing] + median_size = [1, *median_size]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/data/dataset_fingerprint.py` around lines 291 - 296, The code prepends values to median_spacing and median_size using list addition when num_dims == 2; replace those concatenations with unpacking syntax to improve style and clarity (e.g., use median_spacing = [1.0, *median_spacing] and median_size = [1, *median_size]) while keeping the surrounding logic in the dataset_fingerprint.py block that uses num_dims, median_spacing, median_size, anisotropy_ratio, and is_anisotropic unchanged.medicai/trainer/nnunet/utils/io.py (2)
101-116: Note:get_case_iddiffers fromnormalize_case_idin cross_validation.py.This function only strips file extensions, while
normalize_case_idadditionally removes trailing_\d{4}modality suffixes. This is likely intentional (file identification vs. case grouping), but worth documenting to avoid confusion.Consider adding a docstring clarifying the difference
def get_case_id(path): + """ + Extract case identifier from a file path by removing the extension. + + Note: Unlike normalize_case_id() in cross_validation.py, this does NOT + strip modality suffixes (e.g., _0000). Use this for file identification, + use normalize_case_id for grouping multi-modal files by case. + """ name = Path(path).name🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/utils/io.py` around lines 101 - 116, Add a concise docstring to get_case_id explaining that it only strips file extensions (list the handled suffixes) and does not remove trailing modality patterns like _\d{4}, and reference normalize_case_id in cross_validation.py for the alternative behavior; update the function get_case_id's docstring to state intended purpose (file identification vs. case grouping) and that callers should use normalize_case_id when modality suffix normalization is required.
228-244: LGTM with minor style suggestion.DICOM loading correctly handles PixelSpacing and SliceThickness. The optional import pattern for pydicom is appropriate.
Minor: use tuple unpacking for cleaner concatenation (per linter)
if hasattr(dcm, "SliceThickness") and data.ndim >= 3: base = list(spacing) if spacing is not None else [1.0, 1.0] - spacing = (float(dcm.SliceThickness),) + tuple(base) + spacing = (float(dcm.SliceThickness), *base)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/utils/io.py` around lines 228 - 244, The spacing assembly uses explicit list-to-tuple concatenation; update the block in the DICOM branch (variables: dcm, spacing, base, PixelSpacing, SliceThickness) to use tuple unpacking for clarity and style—compute a two-element base tuple when spacing is present or default to (1.0, 1.0), then set spacing = (float(dcm.SliceThickness), *base) so the resulting spacing remains the same but uses cleaner tuple unpacking syntax.medicai/trainer/nnunet/training/augmentations.py (1)
103-113: Creating Keras layers inside__call__is inefficient.The
keras.layers.RandomRotationinstances are created on every call. This is wasteful and can cause issues with graph tracing/compilation. Consider pre-creating the 2D rotation layer in__init__.Proposed fix
In
__init__, add:# For 2D fallback rotation self.rotate_2d_image = keras.layers.RandomRotation( c.rotation_angle_range, fill_mode="constant" ) self.rotate_2d_label = keras.layers.RandomRotation( c.rotation_angle_range, interpolation="nearest", fill_mode="constant" )Then in
__call__:else: - # Placeholder for 2D custom rotation utilizing Keras CV / core ops layer - rotated = keras.layers.RandomRotation(self.config.rotation_angle_range)( - tf.expand_dims(tensor_dict["image"], 0) - ) + rotated = self.rotate_2d_image(tf.expand_dims(tensor_dict["image"], 0)) tensor_dict["image"] = tf.squeeze(rotated, 0) if label is not None: - rotated_lbl = keras.layers.RandomRotation( - self.config.rotation_angle_range, interpolation="nearest" - )(tf.expand_dims(tensor_dict["label"], 0)) + rotated_lbl = self.rotate_2d_label(tf.expand_dims(tensor_dict["label"], 0)) tensor_dict["label"] = tf.squeeze(rotated_lbl, 0)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/training/augmentations.py` around lines 103 - 113, The code currently constructs keras.layers.RandomRotation inside __call__, which recreates layers on every invocation and can break tracing; move creation into __init__ by adding persistent attributes (e.g., self.rotate_2d_image and self.rotate_2d_label) initialized with keras.layers.RandomRotation(c.rotation_angle_range, fill_mode="constant") and interpolation="nearest" for labels, then in __call__ replace the inline keras.layers.RandomRotation(...) calls with self.rotate_2d_image(tf.expand_dims(tensor_dict["image"],0)) / self.rotate_2d_label(tf.expand_dims(tensor_dict["label"],0)) and tf.squeeze results back into tensor_dict["image"] and tensor_dict["label"].medicai/models/nnunet/blocks.py (1)
231-232: Replace Unicode multiplication sign with ASCIIxfor consistency.The docstring uses
×(Unicode multiplication sign) which Ruff flags as ambiguous. Usexor*for cross-platform compatibility.Proposed fix
class SegmentationHead(keras.Layer): - """1×1 (or 1×1×1) convolution → softmax output.""" + """1x1 (or 1x1x1) convolution -> softmax output."""🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/models/nnunet/blocks.py` around lines 231 - 232, The docstring for class SegmentationHead uses a Unicode multiplication sign (×); update that docstring to use an ASCII "x" (or "*" if preferred) — locate the SegmentationHead class and replace "1×1 (or 1×1×1)" with "1x1 (or 1x1x1)" to avoid the ambiguous Unicode character.medicai/trainer/nnunet/data/preprocessing.py (2)
511-521: Remove or implement thenum_workersparameter.The
num_workersparameter is accepted but immediately deleted (Line 521). Either implement parallel preprocessing or remove the parameter from the signature to avoid misleading callers.Proposed fix (remove unused parameter)
def preprocess_dataset( fingerprint, plan, output_dir, configuration="3d_fullres", - num_workers=1, max_cases=None, manifest_file=None, ensure_channel_last=True, ): - del num_workers🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/data/preprocessing.py` around lines 511 - 521, The function preprocess_dataset currently accepts num_workers but immediately deletes it; remove num_workers from the preprocess_dataset signature and delete the del num_workers line, and then update any callers/tests/docs that pass num_workers to call the new signature (or reintroduce a properly implemented parallel processing path using num_workers in preprocess_dataset if parallelism is desired).
107-109: Addstrict=Trueto zip for defensive bounds checking.While
bboxandshapeshould always have matching lengths, addingstrict=Truewould catch any future regressions immediately.Proposed fix
padded_bbox = tuple( - slice(max(0, s.start - pad), min(sh, s.stop + pad)) for s, sh in zip(bbox, shape) + slice(max(0, s.start - pad), min(sh, s.stop + pad)) for s, sh in zip(bbox, shape, strict=True) )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/data/preprocessing.py` around lines 107 - 109, The generator that builds padded_bbox uses zip(bbox, shape) without strict checking; change it to zip(bbox, shape, strict=True) in the padded_bbox assignment so mismatched lengths raise immediately (locate the padded_bbox = tuple(... ) expression that iterates over s, sh in zip(bbox, shape) and add strict=True).medicai/trainer/nnunet/planning/planners.py (2)
715-736: Placeholder subclasses should be documented or removed.
nnUNetPlannerResEncMandnnUNetPlannerResEncLare empty stubs withpassstatements. Either add TODO comments explaining the intended functionality, or remove them until actually implemented to avoid confusion.Option 1: Add TODO with explanation
class nnUNetPlannerResEncM(nnUNetPlanner): """ Planner configured for Medium-scale Residual Encoders. Follows nnU-Net official naming convention for scaling up standard configurations. """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - # Architecture modifications for ResEncM can be overridden here - pass + # TODO: Implement ResEncM-specific architecture modifications + # (increased base_filters, different max_filters, etc.)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/planning/planners.py` around lines 715 - 736, The two empty planner subclasses nnUNetPlannerResEncM and nnUNetPlannerResEncL (their __init__ methods) are placeholder stubs and should not remain ambiguous; either remove these classes entirely if they are not used, or add a clear TODO comment and brief docstring note inside each class (and/or inside the __init__) describing the intended architectural modifications, expected parameters, and a link or issue ID for the future implementation so readers know why they exist and what to implement next.
482-485: Clarify the purpose of storingkernel_sizein NetworkConfig.
kernel_sizeis set toconv_kernels[0]as a "representative" value, but the actual model indynamic_unet.pyderives per-stage kernels independently frompool_op_kernel_sizes. Consider either:
- Removing
kernel_sizefromNetworkConfigif it's unused- Storing the full
conv_kernelslist if per-stage kernels are needed- Documenting that this is for informational purposes only
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/planning/planners.py` around lines 482 - 485, The current code computes per-stage conv_kernels via compute_anisotropic_kernel_sizes and then assigns kernel_size = conv_kernels[0] before placing it into NetworkConfig, but dynamic_unet.py actually derives per-stage kernels from pool_op_kernel_sizes; update NetworkConfig usage to avoid confusion by either (a) removing kernel_size from NetworkConfig entirely if it's unused, (b) storing the full conv_kernels list (replace kernel_size with conv_kernels) so per-stage kernel sizes are available to downstream code, or (c) add a clear docstring/comment where kernel_size is set (in planners.py near compute_anisotropic_kernel_sizes and where NetworkConfig is constructed) stating it is only an informational representative value and not used for model construction; locate references to kernel_size, conv_kernels, NetworkConfig, and dynamic_unet.py/pool_op_kernel_sizes when making the change.medicai/trainer/nnunet/data/normalization.py (1)
78-103: Consider documenting the all-zero foreground edge case behavior.When
nonzero_only=Trueand all voxels are zero, the method returns the image unchanged (Line 94). This is a reasonable fallback, but downstream code should be aware that the returned image won't be normalized. Consider adding a log warning or documenting this in the docstring.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/data/normalization.py` around lines 78 - 103, Update the normalize method to document and/or signal the all-zero-foreground edge case: when nonzero_only is True and fg.size == 0 the function currently returns the image unchanged, so add a brief docstring sentence describing this behavior and either emit a warning via the project's logger (or raise a configurable flag) inside normalize to make downstream code aware; refer to the normalize method and the nonzero_only branch to implement the docstring change and the optional process-aware logging where fg.size == 0.medicai/trainer/nnunet/utils/config.py (1)
14-28: Renamemin/maxparameters to avoid shadowing Python builtins.Using
minandmaxas parameter names shadows the built-in functions, which can cause subtle bugs if those builtins are needed within the method or by readers of the code.Proposed fix
def __init__( self, mean=0.0, std=1.0, - min=0.0, - max=1.0, + min_val=0.0, + max_val=1.0, percentile_00_5=0.0, percentile_99_5=1.0, ): self.mean = mean self.std = std - self.min = min - self.max = max + self.min = min_val + self.max = max_val self.percentile_00_5 = percentile_00_5 self.percentile_99_5 = percentile_99_5Note: The attribute names (
self.min,self.max) can stay the same for backward compatibility with serialized JSON files; only the parameter names need to change.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/utils/config.py` around lines 14 - 28, Rename the __init__ parameters that shadow builtins: change the parameters named min and max to non-shadowing names (e.g., min_val and max_val) in the __init__ signature of the config class, and assign them to the existing attributes self.min and self.max (leave attribute names unchanged for JSON compatibility); update any internal references in __init__ that use the old parameter names (no other callers should need changes if they pass positional args, but prefer keyword 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 `@medicai/models/nnunet/blocks.py`:
- Around line 173-175: The DownBlock.get_config currently returns only
super().get_config(), losing initialization parameters; update
DownBlock.get_config to call base = super().get_config() then extend base with
the DownBlock init args (filters, kernel_size, pool_kernel, spatial_dims,
negative_slope) so they are included in the serialized config; use the exact
parameter names as keys and their current attribute values (e.g., self.filters,
self.kernel_size, self.pool_kernel, self.spatial_dims, self.negative_slope) and
return the merged dict so model.save/from_config roundtrips correctly.
- Around line 122-126: DoubleConvBlock.get_config currently inlines
conv1.get_config() and omits the class __init__ args, breaking Keras
serialization; update DoubleConvBlock.get_config to return a dict that includes
its constructor parameters (filters, kernel_size, spatial_dims, negative_slope)
and nested layer configs for conv1 and conv2 (e.g., store conv1_config =
self.conv1.get_config(), conv2_config = self.conv2.get_config()) instead of
merging them into the top-level config, and ensure super().get_config() is
updated with these keys so that deserialize/clone_model can reconstruct the
block.
- Around line 223-225: UpBlock's get_config currently returns only
super().get_config(), losing the constructor parameters needed for
serialization; update UpBlock.get_config to call base = super().get_config(),
then merge in all UpBlock __init__ arguments (e.g., any parameters like
in_channels, out_channels, kernel_size, stride, use_batchnorm, etc. — match the
exact parameter names from UpBlock.__init__) and return the combined dict so the
layer can be reconstructed from its config.
In `@medicai/models/nnunet/unet.py`:
- Around line 210-222: The call() currently mutates each ResizingND layer's
target_shape (resizers[i].target_shape = target_shape) which is fragile for
compiled/XLA and serialization; instead remove mutation and perform dynamic
resizing inside call() using a functional resize (e.g., keras.ops.image.resize
or tf.image.resize) on seg_outputs[i] to the computed target_shape before
placing into out_dict, leaving the pre-instantiated self.resizers unused for
dynamic paths (or keep them for fixed-size builds only); update the
deep_supervision branch in call()/forward where seg_outputs and resizers are
used so it computes target_shape = ops.shape(seg_outputs[0])[1:-1] and applies
an ops.image.resize call for each auxiliary output rather than assigning to
resizer.target_shape.
- Around line 160-168: The ResizingND instance is being constructed without the
required target_shape or scale_factor, causing a ValueError; update the
ResizingND creation in the block that builds self.resizers (the branch creating
aux_resizer_{stage}) to pass a harmless placeholder scale_factor (e.g.,
scale_factor=1.0) so the constructor succeeds, while preserving the existing
logic that overrides/sets the actual dynamic target later (so keep the name
f"aux_resizer_{stage}" and appending to self.resizers as before).
In `@medicai/trainer/nnunet/data/resampling.py`:
- Around line 160-167: compute_new_shape duplicates the logic of
compute_new_shape_after_resample in planners.py and uses redundant
int(round(...)); extract a single shared utility (e.g., compute_resampled_shape
or compute_new_shape_shared) that takes original_shape, original_spacing,
target_spacing, uses round(...) (no int()), and returns a consistent type (pick
List[int] or tuple and update both callers). Replace compute_new_shape and
compute_new_shape_after_resample bodies to call the new shared function and
update imports/return types so both modules use the same implementation and
type.
In `@medicai/trainer/nnunet/planning/planners.py`:
- Line 371: The file contains a late `from __future__ import annotations` import
that must be moved to the top of the module; open
medicai/trainer/nnunet/planning/planners.py, remove the `from __future__ import
annotations` at its current location (line shown in diff) and add a single `from
__future__ import annotations` immediately at the top of the file just after the
module docstring (before any other imports or code) so the future import is the
first non-docstring statement.
In `@medicai/trainer/nnunet/training/metrics.py`:
- Around line 120-127: The update_state method in metrics.py only unwraps
list/tuple deep-supervision outputs but not dicts; modify update_state (method
name: update_state) to detect when y_pred is a dict and select the primary
output (e.g., y_pred["final"] or the key "final") before computing score, so
mean_dice receives the actual prediction tensor; ensure you still handle
list/tuple as before and keep behavior unchanged when y_pred is already a
tensor.
- Around line 155-162: The update_state in PerClassDiceMetric fails to handle
dict-style y_pred from deep supervision; modify PerClassDiceMetric.update_state
to mirror the fix used elsewhere: if y_pred is a (list, tuple) extract
y_pred[0], and also check if y_pred is a dict (e.g., isinstance(y_pred, dict))
then select the primary output (e.g., y_pred.get("out") or
next(iter(y_pred.values()))) before calling dice_coefficient, then compute
per_class = dice_coefficient(y_true, y_pred, n_classes=self.n_classes) and
update self.dice_sums and self.count as before so dict outputs are handled
safely.
In `@medicai/trainer/nnunet/utils/io.py`:
- Around line 119-130: The infer_spatial_dims function can misclassify thin 3D
volumes (e.g., shape (3,256,256)) as 2D because it treats small first-dimension
sizes (data.shape[0] <= 4) as channels; update infer_spatial_dims to either
accept an optional hint parameter (e.g., is_3d or force_3d) to override the
heuristic for ambiguous cases OR add a clear explanatory comment near the
heuristic branch (the checks using data.shape[0] and data.shape[-1]) documenting
that small-depth volumes may be treated as 2D and recommending callers pass the
hint when they know the data represents a thin 3D volume; ensure the new
parameter is used when provided and falls back to existing logic otherwise.
---
Nitpick comments:
In `@medicai/models/nnunet/blocks.py`:
- Around line 231-232: The docstring for class SegmentationHead uses a Unicode
multiplication sign (×); update that docstring to use an ASCII "x" (or "*" if
preferred) — locate the SegmentationHead class and replace "1×1 (or 1×1×1)" with
"1x1 (or 1x1x1)" to avoid the ambiguous Unicode character.
In `@medicai/models/nnunet/unet.py`:
- Around line 258-263: The for-loop that builds the filter schedule uses an
unused loop variable `i`; update the loop to use a throwaway variable (e.g.,
`for _ in range(n_pooling + 1):`) so linters won't flag an unused variable while
preserving the logic that appends min(f, max_filters) and doubles `f`; adjust
the loop in the schedule construction that references `base_filters`,
`n_pooling`, `max_filters`, `schedule`, and `f` accordingly.
In `@medicai/trainer/nnunet/data/dataset_fingerprint.py`:
- Around line 105-127: The loop using zip(class_ids, label_paths) in
dataset_fingerprint.py can silently truncate when target_class_ids length
differs from label_paths; modify DatasetFingerprint logic to validate lengths
before iterating (e.g., check len(class_ids) == len(label_paths) when
target_class_ids is provided and raise a clear error) or explicitly use
strict=True in the zip to force a ValueError on mismatch; update the code around
class_ids, label_paths and the for ... in zip(...) loop in the function where
counts and total_voxels are computed (references: class_ids, label_paths,
counts, total_voxels) and add a short comment explaining the chosen behavior.
- Around line 291-296: The code prepends values to median_spacing and
median_size using list addition when num_dims == 2; replace those concatenations
with unpacking syntax to improve style and clarity (e.g., use median_spacing =
[1.0, *median_spacing] and median_size = [1, *median_size]) while keeping the
surrounding logic in the dataset_fingerprint.py block that uses num_dims,
median_spacing, median_size, anisotropy_ratio, and is_anisotropic unchanged.
In `@medicai/trainer/nnunet/data/normalization.py`:
- Around line 78-103: Update the normalize method to document and/or signal the
all-zero-foreground edge case: when nonzero_only is True and fg.size == 0 the
function currently returns the image unchanged, so add a brief docstring
sentence describing this behavior and either emit a warning via the project's
logger (or raise a configurable flag) inside normalize to make downstream code
aware; refer to the normalize method and the nonzero_only branch to implement
the docstring change and the optional process-aware logging where fg.size == 0.
In `@medicai/trainer/nnunet/data/preprocessing.py`:
- Around line 511-521: The function preprocess_dataset currently accepts
num_workers but immediately deletes it; remove num_workers from the
preprocess_dataset signature and delete the del num_workers line, and then
update any callers/tests/docs that pass num_workers to call the new signature
(or reintroduce a properly implemented parallel processing path using
num_workers in preprocess_dataset if parallelism is desired).
- Around line 107-109: The generator that builds padded_bbox uses zip(bbox,
shape) without strict checking; change it to zip(bbox, shape, strict=True) in
the padded_bbox assignment so mismatched lengths raise immediately (locate the
padded_bbox = tuple(... ) expression that iterates over s, sh in zip(bbox,
shape) and add strict=True).
In `@medicai/trainer/nnunet/data/resampling.py`:
- Around line 33-36: The length assertion is good but we should make the
resampling ratio computation defensive by using zip(..., strict=True) to satisfy
Ruff B905; update the return in resampling.py (the tuple comprehension that
currently uses zip(original_spacing, target_spacing)) to call
zip(original_spacing, target_spacing, strict=True) so any length mismatch raises
immediately in zip as well (keep the existing assert in place).
In `@medicai/trainer/nnunet/pipeline.py`:
- Around line 56-68: The GPU auto-detection block silently swallows exceptions;
modify the except block to log the failure (e.g., using a module logger obtained
via logging.getLogger(__name__)) so failures in subprocess.run/nvidia-smi
parsing are visible for debugging. Specifically, in the try/except around
subprocess.run/result/memories, catch Exception as e and call logger.debug or
logger.warning with a clear message like "Failed to detect GPU memory via
nvidia-smi" including the exception details and optionally result.stdout, then
fall back to returning float(default) as before.
In `@medicai/trainer/nnunet/planning/planners.py`:
- Around line 715-736: The two empty planner subclasses nnUNetPlannerResEncM and
nnUNetPlannerResEncL (their __init__ methods) are placeholder stubs and should
not remain ambiguous; either remove these classes entirely if they are not used,
or add a clear TODO comment and brief docstring note inside each class (and/or
inside the __init__) describing the intended architectural modifications,
expected parameters, and a link or issue ID for the future implementation so
readers know why they exist and what to implement next.
- Around line 482-485: The current code computes per-stage conv_kernels via
compute_anisotropic_kernel_sizes and then assigns kernel_size = conv_kernels[0]
before placing it into NetworkConfig, but dynamic_unet.py actually derives
per-stage kernels from pool_op_kernel_sizes; update NetworkConfig usage to avoid
confusion by either (a) removing kernel_size from NetworkConfig entirely if it's
unused, (b) storing the full conv_kernels list (replace kernel_size with
conv_kernels) so per-stage kernel sizes are available to downstream code, or (c)
add a clear docstring/comment where kernel_size is set (in planners.py near
compute_anisotropic_kernel_sizes and where NetworkConfig is constructed) stating
it is only an informational representative value and not used for model
construction; locate references to kernel_size, conv_kernels, NetworkConfig, and
dynamic_unet.py/pool_op_kernel_sizes when making the change.
In `@medicai/trainer/nnunet/README.md`:
- Around line 68-75: The "Spacing Rules" section is slightly repetitive and can
be reworded for smoother flow; update the README paragraph(s) under the "Spacing
Rules" heading to vary sentence structure and improve readability by combining
related points (e.g., state that an item’s explicit `spacing` is used first,
then fall back to file metadata for NIfTI/DICOM, and finally to isotropic
defaults), and present the default isotropic values clearly (2D: [1.0, 1.0], 3D:
[1.0, 1.0, 1.0]) while keeping references to `spacing`, NIfTI, DICOM, and
MedicAI intact.
In `@medicai/trainer/nnunet/training/augmentations.py`:
- Around line 103-113: The code currently constructs keras.layers.RandomRotation
inside __call__, which recreates layers on every invocation and can break
tracing; move creation into __init__ by adding persistent attributes (e.g.,
self.rotate_2d_image and self.rotate_2d_label) initialized with
keras.layers.RandomRotation(c.rotation_angle_range, fill_mode="constant") and
interpolation="nearest" for labels, then in __call__ replace the inline
keras.layers.RandomRotation(...) calls with
self.rotate_2d_image(tf.expand_dims(tensor_dict["image"],0)) /
self.rotate_2d_label(tf.expand_dims(tensor_dict["label"],0)) and tf.squeeze
results back into tensor_dict["image"] and tensor_dict["label"].
In `@medicai/trainer/nnunet/utils/config.py`:
- Around line 14-28: Rename the __init__ parameters that shadow builtins: change
the parameters named min and max to non-shadowing names (e.g., min_val and
max_val) in the __init__ signature of the config class, and assign them to the
existing attributes self.min and self.max (leave attribute names unchanged for
JSON compatibility); update any internal references in __init__ that use the old
parameter names (no other callers should need changes if they pass positional
args, but prefer keyword usage).
In `@medicai/trainer/nnunet/utils/io.py`:
- Around line 101-116: Add a concise docstring to get_case_id explaining that it
only strips file extensions (list the handled suffixes) and does not remove
trailing modality patterns like _\d{4}, and reference normalize_case_id in
cross_validation.py for the alternative behavior; update the function
get_case_id's docstring to state intended purpose (file identification vs. case
grouping) and that callers should use normalize_case_id when modality suffix
normalization is required.
- Around line 228-244: The spacing assembly uses explicit list-to-tuple
concatenation; update the block in the DICOM branch (variables: dcm, spacing,
base, PixelSpacing, SliceThickness) to use tuple unpacking for clarity and
style—compute a two-element base tuple when spacing is present or default to
(1.0, 1.0), then set spacing = (float(dcm.SliceThickness), *base) so the
resulting spacing remains the same but uses cleaner tuple unpacking syntax.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 4685c43f-3ba7-4255-bd59-3a8982b3692a
📒 Files selected for processing (26)
medicai/models/nnunet/__init__.pymedicai/models/nnunet/blocks.pymedicai/models/nnunet/dynamic_unet.pymedicai/models/nnunet/unet.pymedicai/trainer/__init__.pymedicai/trainer/nnunet/README.mdmedicai/trainer/nnunet/ROADMAP.mdmedicai/trainer/nnunet/__init__.pymedicai/trainer/nnunet/cross_validation.pymedicai/trainer/nnunet/data/__init__.pymedicai/trainer/nnunet/data/dataset_fingerprint.pymedicai/trainer/nnunet/data/manifest.pymedicai/trainer/nnunet/data/normalization.pymedicai/trainer/nnunet/data/preprocessing.pymedicai/trainer/nnunet/data/resampling.pymedicai/trainer/nnunet/pipeline.pymedicai/trainer/nnunet/planning/__init__.pymedicai/trainer/nnunet/planning/planners.pymedicai/trainer/nnunet/training/__init__.pymedicai/trainer/nnunet/training/augmentations.pymedicai/trainer/nnunet/training/losses.pymedicai/trainer/nnunet/training/metrics.pymedicai/trainer/nnunet/training/trainer.pymedicai/trainer/nnunet/utils/__init__.pymedicai/trainer/nnunet/utils/config.pymedicai/trainer/nnunet/utils/io.py
Owner
Author
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a Keras 3 implementation of the nnU-Net framework, providing a self-configuring pipeline for medical image segmentation. The contribution includes a dynamic U-Net architecture with deep supervision, automated planning based on dataset fingerprints, and a flexible manifest-based data handling system. The review feedback suggests optimizing memory usage during model weight initialization by using model.build instead of dummy tensors, implementing parallel execution for preprocessing by utilizing the currently ignored num_workers parameter, and making the gradient clipping threshold a configurable parameter in the training configuration.
| n_mod = net_cfg.n_modalities if net_cfg else 1 | ||
| n_outputs = net_cfg.n_classes if net_cfg else 2 | ||
|
|
||
| dummy = np.zeros([1] + patch_size + [n_mod], dtype=np.float32) |
There was a problem hiding this comment.
Creating a dummy tensor using np.zeros and then passing it to the model for weight initialization is a common pattern, but it can be memory-intensive for large 3D patches. Consider using model.build(input_shape) if the model supports it, to avoid unnecessary memory allocation.
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (12)
medicai/models/nnunet/unet.py (2)
148-169:self.resizerslist is now unused dead code.The
resizerslist is populated during__init__but is never referenced incall()since the deep supervision branch now usesops.image.resizedirectly. This dead code adds maintenance burden and memory overhead.♻️ Remove unused resizers
# ---- Segmentation heads (one per decoder level for deep supervision) self.seg_heads = [] - self.resizers = [] for stage in range(n_pooling): head = SegmentationHead( n_classes=n_classes, spatial_dims=spatial_dims, activation=output_activation, name=f"seg_head_{stage}", ) self.seg_heads.append(head) - - # Resizers for deep supervision (aux outputs → full resolution) - if stage > 0: - resizer = ResizingND( - scale_factor=1.0, # placeholder; dynamic resize in call() - interpolation="bilinear" if spatial_dims == 2 else "trilinear", - name=f"aux_resizer_{stage}", - ) - self.resizers.append(resizer) - else: - self.resizers.append(None)Also remove the
ResizingNDimport on line 42 if no longer needed elsewhere.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/models/nnunet/unet.py` around lines 148 - 169, Remove the unused self.resizers list and its population in __init__ since call() now uses ops.image.resize; delete the self.resizers attribute, the branch that appends ResizingND or None inside the for loop that creates seg_heads (symbols: self.resizers, SegmentationHead, aux_resizer_{stage}), and any related members, and also remove the ResizingND import if it is not used elsewhere in the file; ensure only seg_heads remain created in the loop and run tests to confirm no other references to self.resizers or ResizingND exist.
32-42: Unused import after refactoring.
ResizingNDis imported but now only used to build the unusedself.resizerslist. If the resizers are removed as suggested, this import becomes dead code.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/models/nnunet/unet.py` around lines 32 - 42, ResizingND is no longer used because the now-unused attribute self.resizers was removed during refactoring; remove the dead import from the top-level imports and any leftover references to ResizingND (search for the symbol ResizingND and the attribute self.resizers in the UNet class and __init__ or any methods) so the module has no unused imports, and run a quick lint to confirm no remaining references.medicai/trainer/nnunet/utils/config.py (2)
109-112:from_jsonmethods are fragile to extra/unknown keys in JSON.Using
cls(**data)directly will raiseTypeErrorif the JSON contains keys not in the constructor signature. This makes schema evolution difficult and can cause failures when loading files saved by newer versions.Consider filtering to known parameters or using a more defensive unpacking approach.
♻️ Defensive unpacking example
`@classmethod` def from_json(cls, path): data = load_json(path) import inspect valid_params = set(inspect.signature(cls.__init__).parameters.keys()) - {'self'} filtered = {k: v for k, v in data.items() if k in valid_params} return cls(**filtered)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/utils/config.py` around lines 109 - 112, The from_json method is fragile to extra keys because it passes raw JSON into cls(**data); update from_json (method name) to defensively filter loaded data (from load_json) to only constructor-accepted parameters before calling cls(**filtered). Use inspect.signature(cls.__init__) to derive valid parameter names (exclude 'self') or, for dataclasses, read dataclass fields, then build a filtered dict of {k:v for k,v in data.items() if k in valid_params} and return cls(**filtered).
335-346:TrainingConfig.from_jsonhandles legacy keys but remains vulnerable to unknown keys.While legacy keys are explicitly removed, any other unexpected key in the JSON will still cause a
TypeError. The pattern could be extended to filter to known parameters.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/utils/config.py` around lines 335 - 346, TrainingConfig.from_json currently removes a few legacy keys but will still raise TypeError if the JSON contains any unknown keys; update from_json to whitelist only the known constructor parameters before calling cls(**data): load the JSON via load_json(path), remove legacy keys as done, then compute the allowed keys from the class constructor (e.g., inspect.signature(cls.__init__).parameters or dataclasses.fields(cls) if TrainingConfig is a dataclass) and filter data to only those keys, then call cls(**filtered_data) instead of cls(**data) to avoid unexpected-key TypeErrors.medicai/trainer/nnunet/data/dataset_fingerprint.py (1)
265-268: Consider addingstrict=Truetozipfor consistency.While
np.unique(..., return_counts=True)guaranteesuniqueandcountshave the same length, addingstrict=Truefor consistency with line 106 and to be defensive against future refactoring would be beneficial.♻️ Proposed fix
- unique, counts = np.unique(label_data.astype(np.int64), return_counts=True) - for val, count in zip(unique, counts): + unique, counts = np.unique(label_data.astype(np.int64), return_counts=True) + for val, count in zip(unique, counts, strict=True):🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/data/dataset_fingerprint.py` around lines 265 - 268, In the loop that aggregates voxel counts (using np.unique on label_data and updating class_voxel_counts), make the zip call defensive by adding strict=True (i.e., change zip(unique, counts) to zip(unique, counts, strict=True)) so mismatched lengths raise immediately and remain consistent with the other occurrence on line 106; this ensures the pairing of values and counts is enforced when iterating and helps catch future refactors that might break the assumption.medicai/trainer/nnunet/training/augmentations.py (2)
101-107: Redundant 2D/3D conditional for rotation.Both branches of the
if is_3d/elseconditional execute identical code. This can be simplified.♻️ Simplify redundant conditional
# 3. Random Rotation - if is_3d: - tensor_dict = self.rotate(tensor_dict).data - else: - # 2D rotation using the same RandRotate with consistent - # randomness for both image and label - tensor_dict = self.rotate(tensor_dict).data + tensor_dict = self.rotate(tensor_dict).data🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/training/augmentations.py` around lines 101 - 107, The if/else around is_3d is redundant because both branches call self.rotate(tensor_dict).data; remove the conditional and replace it with a single call assigning tensor_dict = self.rotate(tensor_dict).data (keep the is_3d symbol if needed elsewhere), updating any surrounding comments to reflect one unified rotation step; target the block that currently references is_3d, tensor_dict and the rotate method in this function.
59-65:flip3applied unconditionally to 3D but always created.
self.flip3withspatial_axis=[2]is created for all pipelines but only used for 3D inputs. For 2D-only workflows, this is wasted memory.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/training/augmentations.py` around lines 59 - 65, The code unconditionally instantiates self.flip3 = RandFlip(keys=keys, prob=c.p_mirror, spatial_axis=[2]) even for 2D pipelines; guard creation so RandFlip for spatial_axis=[2] is only created when the pipeline is 3D (e.g., when a config flag like c.spatial_dims or c.is_3d indicates 3D) — otherwise set self.flip3 = None (or omit it) and ensure downstream code checks for its existence before use; update the instantiation around self.flip, self.flip2, self.flip3 to conditionally create self.flip3 based on that 3D indicator while keeping RandRotate (self.rotate) unchanged.medicai/trainer/nnunet/data/manifest.py (1)
133-135: Confusingor Noneexpression with list comprehension.The expression
[[int(v) for v in region] for region in item_dict.get("regions", [])] or Noneevaluates the entire list comprehension first, then appliesor None. An empty list[]is falsy, so this works, but the precedence and intent are unclear at first glance.♻️ Clearer expression
- regions=[[int(v) for v in region] for region in item_dict.get("regions", [])] - or None, + regions=[[int(v) for v in region] for region in item_dict.get("regions", [])] or None,Or for even clearer intent:
raw_regions = item_dict.get("regions", []) regions = [[int(v) for v in region] for region in raw_regions] if raw_regions else None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/data/manifest.py` around lines 133 - 135, The regions assignment uses a list comprehension followed by "or None", which is confusing; change it to first fetch raw_regions = item_dict.get("regions", []) and then set regions = [[int(v) for v in region] for region in raw_regions] if raw_regions else None so intent and short-circuiting are explicit (refer to the regions assignment that reads item_dict.get("regions", []) and the surrounding manifest parsing code).medicai/trainer/nnunet/pipeline.py (1)
466-468: Consider using iterable unpacking for clarity.Per static analysis suggestion, this can be slightly cleaner with unpacking.
♻️ Use iterable unpacking
- dummy = np.zeros([1] + patch_size + [n_mod], dtype=np.float32) + dummy = np.zeros([1, *patch_size, n_mod], dtype=np.float32)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/pipeline.py` around lines 466 - 468, Replace the concatenated list construction for the dummy array shape with iterable unpacking for clarity: change the np.zeros([1] + patch_size + [n_mod], dtype=np.float32) call to use np.zeros([1, *patch_size, n_mod], dtype=np.float32) while leaving the subsequent model(dummy, training=False) and model.load_weights(str(model_weights_path)) lines unchanged.medicai/trainer/nnunet/planning/planners.py (2)
715-736: Placeholder subclasses with TODOs.These classes are noted as placeholders for future ResEnc variants. The TODOs clearly document the intended extensions.
Would you like me to help implement the ResEncM/ResEncL-specific architecture modifications, or open issues to track this work?
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/planning/planners.py` around lines 715 - 736, The two planner subclasses nnUNetPlannerResEncM and nnUNetPlannerResEncL are left as placeholders with TODOs in their __init__ methods; implement ResEnc-specific configuration by overriding __init__ in each to call super().__init__(...) then set planner attributes (e.g., base_filters, max_filters, num_decoder_heads, residual_encoder_flag) or call a helper like configure_resenc(scale='M'|'L') to apply the scale-specific defaults; ensure any new attributes are consistent with the parent nnUNetPlanner expected fields and add minimal validation (raise ValueError on invalid values) so the planners behave like concrete configurations rather than stubs.
159-177: Confusing variable naming in list comprehension.Line 174 iterates
for iso in is_aniso, butis_anisomeans "is anisotropic", makingisoa misleading name. The logic is correct (kernel=1 for anisotropic axes), but readability would improve with clearer naming.Suggested clarity improvement
- kernel = [1 if iso else 3 for iso in is_aniso] + kernel = [1 if aniso else 3 for aniso in is_aniso]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/planning/planners.py` around lines 159 - 177, The list comprehension in compute_anisotropic_kernel_sizes uses the misleading loop variable name "iso" (for example: "for iso in is_aniso"); change the names to be clearer by renaming the boolean list to something like "anisotropic_axes" or "is_aniso_axis" and the loop variable to "is_aniso_axis" or "anisotropic" so the comprehension becomes e.g. "[1 if is_aniso_axis else 3 for is_aniso_axis in anisotropic_axes]"; update only the variable names in that scope to improve readability while keeping the existing logic.medicai/trainer/nnunet/data/preprocessing.py (1)
207-353: Consider refactoring to reduce code duplication.The
_load_labelsfunction has ~150 lines with repeated patterns for loading and normalizing label data across different branches (regions, single multi-class, multiple binary files, standard multi-class). Extracting a helper like_load_and_normalize_single_label(path, ...)could reduce duplication and improve maintainability.This is not blocking but would improve long-term code health.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/data/preprocessing.py` around lines 207 - 353, The _load_labels function contains repeated sequences for loading and preparing a single label (load_medical_image -> determine candidate_spacing -> infer_spatial_dims -> normalize_layout -> collapse_single_channel -> squeeze channel if needed), so extract that sequence into a helper (e.g., _load_and_normalize_single_label(label_path, spatial_dims, ensure_channel_last, original_spacing_override=None, label_layout=None)) that returns (label_data, label_dims, label_spacing) or just normalized label_data; then replace the duplicated blocks in _load_labels (the regions branch, the single-file multi-label branch, the per-file loop for multi-label lists, and the default single-path branch) to call this helper and continue with the per-branch logic (building region masks, class-specific channels, binary mask handling, setting ignored ids). Ensure the helper preserves channel-last semantics and collapsing behavior so downstream code (e.g., computing class_ids, np.isin, stacking) remains unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@medicai/trainer/nnunet/data/dataset_fingerprint.py`:
- Around line 245-248: The loop that aggregates region-based multi-label counts
drops the final region because _collect_multilabel_class_stats yields class IDs
starting at 1 while class_voxel_counts is keyed 0..n-1; change the aggregation
in the block referencing multilabel_counts, class_voxel_counts and class_names
to map region IDs to zero-based indices (e.g. idx = class_id - 1) and guard with
0 <= idx < len(class_names) before doing class_voxel_counts[idx] += count so the
last region (when class_id == len(class_names)) is included.
In `@medicai/trainer/nnunet/data/manifest.py`:
- Around line 1-9: The module uses typing names (List, Optional, Union, Dict,
Any) in dataclass/class signatures but never imports them, causing NameError;
fix by adding the missing imports from typing (import List, Optional, Union,
Dict, Any) at the top of manifest.py so symbols used in the dataclass and any
class/type annotations resolve correctly (look for uses in the dataclass and
functions referencing List/Optional/Union/Dict/Any).
In `@medicai/trainer/nnunet/data/preprocessing.py`:
- Around line 1-7: The module uses type annotations like Tuple, List, Dict, and
Optional but does not import them, which will raise NameError during runtime
introspection; update the imports to include these symbols from typing (e.g.,
add an import for Tuple, List, Dict, Optional) alongside the existing imports so
functions/classes that reference those types (search for usages of Tuple, List,
Dict, Optional in this file) resolve correctly when using
typing.get_type_hints() or other runtime introspection.
In `@medicai/trainer/nnunet/pipeline.py`:
- Around line 192-201: The import of scipy.ndimage.zoom inside the per-batch
code causes repeated overhead; move the import to module-level (top of
medicai/trainer/nnunet/pipeline.py) or to the class __init__, e.g., from
scipy.ndimage import zoom as ndimage_zoom, and update the __getitem__ usage that
currently calls ndimage_zoom(...) (the block creating ds_label and assigning
y_dict[f"aux_{i}"]) to use that module-level symbol instead; ensure there are no
remaining local imports in the __getitem__ method.
In `@medicai/trainer/nnunet/planning/planners.py`:
- Around line 17-21: The module is missing imports for typing names used in
annotations (Sequence, List, Optional, Union); add an import statement such as
"from typing import Sequence, List, Optional, Union" near the top of
medicaI/trainer/nnunet/planning/planners.py so references in functions/classes
that use Sequence, List, Optional, or Union resolve (search for usages in any
function or class like planning-related helpers in this file and ensure the
import is present).
In `@medicai/trainer/nnunet/training/augmentations.py`:
- Around line 84-91: The code mutates transform attributes (self.crop.roi_size /
self.crop.spatial_size) inside __call__ based on patch_size and is_3d, which is
fragile for concurrent/compiled/traced runs; instead construct the crop/resize
transform with the correct spatial size up front or use a stateless resize/crop
operation at call time: replace in-place assignment of self.crop.roi_size and
self.crop.spatial_size by either (a) creating a new crop/resize transform
instance configured from patch_size/is_3d before applying it, or (b) using a
stateless transform API that accepts the target size (e.g., a functional
resize/crop) so that self.crop remains immutable during __call__; update
references to self.crop, patch_size, is_3d, and __call__ accordingly.
In `@medicai/trainer/nnunet/training/metrics.py`:
- Around line 205-212: The hausdorff_distance_95 function currently returns
float("inf") when either mask_true or mask_pred is empty, which lets infinite
values silently corrupt mean/std in aggregate_fold_results; change that sentinel
to np.nan (consistent with the scipy-unavailable branch that returns
float("nan")) so callers computing statistics ignore empty-mask cases, i.e.,
replace the float("inf") return in hausdorff_distance_95 with np.nan and ensure
numpy (np) is available/imported where the function is defined; you may also add
a short comment noting this mirrors the scipy-unavailable behavior.
- Around line 13-19: The module uses Optional[int] in type hints (occurrences at
the Optional[int] annotations) but never imports Optional, causing a NameError;
fix by adding "from typing import Optional" to the import block (alongside
existing imports at top of medicai/trainer/nnunet/training/metrics.py) so the
Optional symbol is defined, or alternatively replace those annotations with "int
| None" if you prefer Python 3.10 union types.
In `@medicai/trainer/nnunet/training/trainer.py`:
- Around line 79-87: The code sets self.net_cfg via net_cfg_map but then
unconditionally accesses self.net_cfg.n_classes which can raise AttributeError
if configuration isn't present or plan.* is None; update the assignment logic
around net_cfg_map / self.net_cfg (and the subsequent self.n_classes
initialization) to first check whether self.net_cfg is truthy and either (a) set
self.n_classes to a safe default (e.g., 2) and log a warning, or (b) raise a
clear exception indicating the unknown configuration; modify the block that
references net_cfg_map, self.net_cfg, and self.n_classes (and optionally
plan.plan_3d_fullres / plan.plan_3d_lowres / plan.plan_2d) to perform this guard
before accessing .n_classes.
---
Nitpick comments:
In `@medicai/models/nnunet/unet.py`:
- Around line 148-169: Remove the unused self.resizers list and its population
in __init__ since call() now uses ops.image.resize; delete the self.resizers
attribute, the branch that appends ResizingND or None inside the for loop that
creates seg_heads (symbols: self.resizers, SegmentationHead,
aux_resizer_{stage}), and any related members, and also remove the ResizingND
import if it is not used elsewhere in the file; ensure only seg_heads remain
created in the loop and run tests to confirm no other references to
self.resizers or ResizingND exist.
- Around line 32-42: ResizingND is no longer used because the now-unused
attribute self.resizers was removed during refactoring; remove the dead import
from the top-level imports and any leftover references to ResizingND (search for
the symbol ResizingND and the attribute self.resizers in the UNet class and
__init__ or any methods) so the module has no unused imports, and run a quick
lint to confirm no remaining references.
In `@medicai/trainer/nnunet/data/dataset_fingerprint.py`:
- Around line 265-268: In the loop that aggregates voxel counts (using np.unique
on label_data and updating class_voxel_counts), make the zip call defensive by
adding strict=True (i.e., change zip(unique, counts) to zip(unique, counts,
strict=True)) so mismatched lengths raise immediately and remain consistent with
the other occurrence on line 106; this ensures the pairing of values and counts
is enforced when iterating and helps catch future refactors that might break the
assumption.
In `@medicai/trainer/nnunet/data/manifest.py`:
- Around line 133-135: The regions assignment uses a list comprehension followed
by "or None", which is confusing; change it to first fetch raw_regions =
item_dict.get("regions", []) and then set regions = [[int(v) for v in region]
for region in raw_regions] if raw_regions else None so intent and
short-circuiting are explicit (refer to the regions assignment that reads
item_dict.get("regions", []) and the surrounding manifest parsing code).
In `@medicai/trainer/nnunet/data/preprocessing.py`:
- Around line 207-353: The _load_labels function contains repeated sequences for
loading and preparing a single label (load_medical_image -> determine
candidate_spacing -> infer_spatial_dims -> normalize_layout ->
collapse_single_channel -> squeeze channel if needed), so extract that sequence
into a helper (e.g., _load_and_normalize_single_label(label_path, spatial_dims,
ensure_channel_last, original_spacing_override=None, label_layout=None)) that
returns (label_data, label_dims, label_spacing) or just normalized label_data;
then replace the duplicated blocks in _load_labels (the regions branch, the
single-file multi-label branch, the per-file loop for multi-label lists, and the
default single-path branch) to call this helper and continue with the per-branch
logic (building region masks, class-specific channels, binary mask handling,
setting ignored ids). Ensure the helper preserves channel-last semantics and
collapsing behavior so downstream code (e.g., computing class_ids, np.isin,
stacking) remains unchanged.
In `@medicai/trainer/nnunet/pipeline.py`:
- Around line 466-468: Replace the concatenated list construction for the dummy
array shape with iterable unpacking for clarity: change the np.zeros([1] +
patch_size + [n_mod], dtype=np.float32) call to use np.zeros([1, *patch_size,
n_mod], dtype=np.float32) while leaving the subsequent model(dummy,
training=False) and model.load_weights(str(model_weights_path)) lines unchanged.
In `@medicai/trainer/nnunet/planning/planners.py`:
- Around line 715-736: The two planner subclasses nnUNetPlannerResEncM and
nnUNetPlannerResEncL are left as placeholders with TODOs in their __init__
methods; implement ResEnc-specific configuration by overriding __init__ in each
to call super().__init__(...) then set planner attributes (e.g., base_filters,
max_filters, num_decoder_heads, residual_encoder_flag) or call a helper like
configure_resenc(scale='M'|'L') to apply the scale-specific defaults; ensure any
new attributes are consistent with the parent nnUNetPlanner expected fields and
add minimal validation (raise ValueError on invalid values) so the planners
behave like concrete configurations rather than stubs.
- Around line 159-177: The list comprehension in
compute_anisotropic_kernel_sizes uses the misleading loop variable name "iso"
(for example: "for iso in is_aniso"); change the names to be clearer by renaming
the boolean list to something like "anisotropic_axes" or "is_aniso_axis" and the
loop variable to "is_aniso_axis" or "anisotropic" so the comprehension becomes
e.g. "[1 if is_aniso_axis else 3 for is_aniso_axis in anisotropic_axes]"; update
only the variable names in that scope to improve readability while keeping the
existing logic.
In `@medicai/trainer/nnunet/training/augmentations.py`:
- Around line 101-107: The if/else around is_3d is redundant because both
branches call self.rotate(tensor_dict).data; remove the conditional and replace
it with a single call assigning tensor_dict = self.rotate(tensor_dict).data
(keep the is_3d symbol if needed elsewhere), updating any surrounding comments
to reflect one unified rotation step; target the block that currently references
is_3d, tensor_dict and the rotate method in this function.
- Around line 59-65: The code unconditionally instantiates self.flip3 =
RandFlip(keys=keys, prob=c.p_mirror, spatial_axis=[2]) even for 2D pipelines;
guard creation so RandFlip for spatial_axis=[2] is only created when the
pipeline is 3D (e.g., when a config flag like c.spatial_dims or c.is_3d
indicates 3D) — otherwise set self.flip3 = None (or omit it) and ensure
downstream code checks for its existence before use; update the instantiation
around self.flip, self.flip2, self.flip3 to conditionally create self.flip3
based on that 3D indicator while keeping RandRotate (self.rotate) unchanged.
In `@medicai/trainer/nnunet/utils/config.py`:
- Around line 109-112: The from_json method is fragile to extra keys because it
passes raw JSON into cls(**data); update from_json (method name) to defensively
filter loaded data (from load_json) to only constructor-accepted parameters
before calling cls(**filtered). Use inspect.signature(cls.__init__) to derive
valid parameter names (exclude 'self') or, for dataclasses, read dataclass
fields, then build a filtered dict of {k:v for k,v in data.items() if k in
valid_params} and return cls(**filtered).
- Around line 335-346: TrainingConfig.from_json currently removes a few legacy
keys but will still raise TypeError if the JSON contains any unknown keys;
update from_json to whitelist only the known constructor parameters before
calling cls(**data): load the JSON via load_json(path), remove legacy keys as
done, then compute the allowed keys from the class constructor (e.g.,
inspect.signature(cls.__init__).parameters or dataclasses.fields(cls) if
TrainingConfig is a dataclass) and filter data to only those keys, then call
cls(**filtered_data) instead of cls(**data) to avoid unexpected-key TypeErrors.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e55fd2a6-e062-4280-a529-27bdad993eb2
📒 Files selected for processing (14)
medicai/models/nnunet/blocks.pymedicai/models/nnunet/unet.pymedicai/trainer/nnunet/data/dataset_fingerprint.pymedicai/trainer/nnunet/data/manifest.pymedicai/trainer/nnunet/data/normalization.pymedicai/trainer/nnunet/data/preprocessing.pymedicai/trainer/nnunet/data/resampling.pymedicai/trainer/nnunet/pipeline.pymedicai/trainer/nnunet/planning/planners.pymedicai/trainer/nnunet/training/augmentations.pymedicai/trainer/nnunet/training/metrics.pymedicai/trainer/nnunet/training/trainer.pymedicai/trainer/nnunet/utils/config.pymedicai/trainer/nnunet/utils/io.py
🚧 Files skipped from review as they are similar to previous changes (3)
- medicai/trainer/nnunet/data/resampling.py
- medicai/models/nnunet/blocks.py
- medicai/trainer/nnunet/utils/io.py
Comment on lines
+245
to
+248
| for class_id, count in multilabel_counts.items(): | ||
| if class_id < len(class_names): | ||
| class_voxel_counts[class_id] += count | ||
| total_voxels += case_total |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check how regions and class_names relate in multi-label scenarios
rg -n "regions" --type py medicai/trainer/nnunet/data/ -A 3 -B 3Repository: innat/medic-ai
Length of output: 9339
🏁 Script executed:
head -280 medicai/trainer/nnunet/data/dataset_fingerprint.py | tail -100Repository: innat/medic-ai
Length of output: 4166
🏁 Script executed:
sed -n '200,260p' medicai/trainer/nnunet/data/dataset_fingerprint.py | cat -nRepository: innat/medic-ai
Length of output: 3092
🏁 Script executed:
sed -n '150,210p' medicai/trainer/nnunet/data/dataset_fingerprint.py | cat -nRepository: innat/medic-ai
Length of output: 2697
🏁 Script executed:
grep -n "_collect_multilabel_class_stats" medicai/trainer/nnunet/data/dataset_fingerprint.pyRepository: innat/medic-ai
Length of output: 179
🏁 Script executed:
sed -n '45,120p' medicai/trainer/nnunet/data/dataset_fingerprint.py | cat -nRepository: innat/medic-ai
Length of output: 3860
🏁 Script executed:
sed -n '115,160p' medicai/trainer/nnunet/data/dataset_fingerprint.py | cat -nRepository: innat/medic-ai
Length of output: 1807
🏁 Script executed:
grep -n "class_names" medicai/trainer/nnunet/data/manifest.py | head -20Repository: innat/medic-ai
Length of output: 292
🏁 Script executed:
sed -n '28,45p' medicai/trainer/nnunet/data/dataset_fingerprint.pyRepository: innat/medic-ai
Length of output: 629
🏁 Script executed:
sed -n '36,42p' medicai/trainer/nnunet/data/dataset_fingerprint.pyRepository: innat/medic-ai
Length of output: 410
Region-based multi-label counts may skip the final region.
When using region-based multi-label, _collect_multilabel_class_stats returns counts with region indices starting at 1 (via enumerate(regions, start=1)). However, class_voxel_counts is initialized with keys range(n_classes) (0 to n_classes-1), and the condition if class_id < len(class_names) will exclude any region index equal to len(class_names). If the number of regions equals the number of class names, the final region's voxel count will be silently dropped.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/trainer/nnunet/data/dataset_fingerprint.py` around lines 245 - 248,
The loop that aggregates region-based multi-label counts drops the final region
because _collect_multilabel_class_stats yields class IDs starting at 1 while
class_voxel_counts is keyed 0..n-1; change the aggregation in the block
referencing multilabel_counts, class_voxel_counts and class_names to map region
IDs to zero-based indices (e.g. idx = class_id - 1) and guard with 0 <= idx <
len(class_names) before doing class_voxel_counts[idx] += count so the last
region (when class_id == len(class_names)) is included.
Comment on lines
+84
to
+91
| if hasattr(self.crop, "roi_size") and patch_size is not None: | ||
| roi_size = list(patch_size) | ||
| if is_3d: | ||
| self.crop.roi_size = roi_size | ||
| else: | ||
| self.crop.roi_size = roi_size[-2:] | ||
| elif hasattr(self.crop, "spatial_size") and patch_size is not None: | ||
| self.crop.spatial_size = tuple(patch_size) |
There was a problem hiding this comment.
Mutating transform attributes at call time is fragile.
Dynamically setting self.crop.roi_size or self.crop.spatial_size during __call__ can cause issues with concurrent batch processing, compiled execution, or functional model tracing (similar to the ResizingND issue in unet.py).
Consider instantiating transforms with the correct sizes upfront or using stateless resize operations.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/trainer/nnunet/training/augmentations.py` around lines 84 - 91, The
code mutates transform attributes (self.crop.roi_size / self.crop.spatial_size)
inside __call__ based on patch_size and is_3d, which is fragile for
concurrent/compiled/traced runs; instead construct the crop/resize transform
with the correct spatial size up front or use a stateless resize/crop operation
at call time: replace in-place assignment of self.crop.roi_size and
self.crop.spatial_size by either (a) creating a new crop/resize transform
instance configured from patch_size/is_3d before applying it, or (b) using a
stateless transform API that accepts the target size (e.g., a functional
resize/crop) so that self.crop remains immutable during __call__; update
references to self.crop, patch_size, is_3d, and __call__ accordingly.
Comment on lines
+205
to
+212
| if not _SCIPY_AVAILABLE: | ||
| return float("nan") | ||
|
|
||
| mask_true = (y_true_np == class_idx).astype(np.uint8) | ||
| mask_pred = (y_pred_np == class_idx).astype(np.uint8) | ||
|
|
||
| if mask_true.sum() == 0 or mask_pred.sum() == 0: | ||
| return float("inf") |
There was a problem hiding this comment.
inf return value may silently corrupt aggregated metrics.
When either mask is empty, hausdorff_distance_95 returns float("inf"). Per cross_validation.py:157-159, aggregate_fold_results computes mean/std directly on arrays that may contain these sentinel values, causing inf to propagate into summary statistics without warning.
Consider returning np.nan instead of inf for empty masks (like the scipy-unavailable case), or document that callers must filter these values before aggregation.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/trainer/nnunet/training/metrics.py` around lines 205 - 212, The
hausdorff_distance_95 function currently returns float("inf") when either
mask_true or mask_pred is empty, which lets infinite values silently corrupt
mean/std in aggregate_fold_results; change that sentinel to np.nan (consistent
with the scipy-unavailable branch that returns float("nan")) so callers
computing statistics ignore empty-mask cases, i.e., replace the float("inf")
return in hausdorff_distance_95 with np.nan and ensure numpy (np) is
available/imported where the function is defined; you may also add a short
comment noting this mirrors the scipy-unavailable behavior.
Owner
Author
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a Keras 3 implementation of the nnU-Net pipeline, including a dynamic UNet architecture and a manifest-based preprocessing system. The review identifies critical bugs in the UNet forward pass regarding output resolution and 3D resizing compatibility, as well as NameErrors in the preprocessing script. Suggestions were also provided to improve the flexibility of augmentations and the efficiency of the data loading pipeline.
Comment on lines
+166
to
+183
| # Full resolution is first decoder output (index 0) | ||
| if training and self.deep_supervision: | ||
| # Return a dictionary of outputs for Keras 3 multi-output training | ||
| # All outputs are resized to match segment_outputs[0] shape | ||
| target_shape = ops.shape(seg_outputs[0])[1:-1] | ||
| out_dict = {"final": seg_outputs[0]} | ||
| for i in range(1, len(seg_outputs)): | ||
| # Functional resize — avoids fragile layer mutation | ||
| interp = "bilinear" if self.spatial_dims == 2 else "trilinear" | ||
| resized = ops.image.resize( | ||
| seg_outputs[i], | ||
| size=target_shape, | ||
| interpolation=interp, | ||
| ) | ||
| out_dict[f"aux_{i-1}"] = resized | ||
| return out_dict | ||
| else: | ||
| return seg_outputs[0] # single full-resolution output |
There was a problem hiding this comment.
The UNet implementation has several critical issues in its call method:
- Output Order Inversion: The decoder outputs are collected from the bottleneck upwards, meaning
seg_outputs[0]is the lowest resolution andseg_outputs[-1]is the full resolution. The code incorrectly treatsseg_outputs[0]as the 'final' output, which will result in the model producing a downsampled segmentation map. - Deep Supervision Logic: Because of the inversion,
target_shapeis set to the smallest resolution, and larger auxiliary maps are downsampled to it. This is the opposite of standard deep supervision where auxiliary maps are compared to downsampled labels at their native resolutions. - 3D Resizing Incompatibility:
ops.image.resizein Keras 3 only supports 2D images (4D tensors). For 3D medical volumes (5D tensors), this operation will fail. Additionally,'trilinear'is not a valid interpolation mode forops.image.resize.
Comment on lines
+594
to
+618
| if manifest_file is None or not Path(manifest_file).exists(): | ||
| raise ValueError( | ||
| "manifest_file must be provided for preprocessing. " | ||
| "Legacy MSD layout is no longer supported." | ||
| ) | ||
|
|
||
| worker_fn = partial( | ||
| _process_item_helper, | ||
| output_dir=output_dir, | ||
| properties_dir=properties_dir, | ||
| manifest=manifest, | ||
| target_spacing=target_spacing, | ||
| plan=plan, | ||
| fingerprint=fingerprint, | ||
| configuration=configuration, | ||
| ensure_channel_last=ensure_channel_last, | ||
| ) | ||
|
|
||
| workers = num_workers if num_workers is not None else max(1, multiprocessing.cpu_count() - 1) | ||
| if workers > 1: | ||
| with concurrent.futures.ProcessPoolExecutor(max_workers=workers) as executor: | ||
| list(tqdm(executor.map(worker_fn, items), total=len(items), desc=f"Preprocessing ({configuration})")) | ||
| else: | ||
| for item in tqdm(items, desc=f"Preprocessing ({configuration})"): | ||
| worker_fn(item) |
There was a problem hiding this comment.
The preprocess_dataset function uses manifest and items variables that are never defined, which will cause a NameError. These should be initialized by loading the manifest from the provided manifest_file.
Comment on lines
+59
to
+61
| self.flip = RandFlip(keys=keys, prob=c.p_mirror, spatial_axis=[0]) | ||
| self.flip2 = RandFlip(keys=keys, prob=c.p_mirror, spatial_axis=[1]) | ||
| self.flip3 = RandFlip(keys=keys, prob=c.p_mirror, spatial_axis=[2]) |
There was a problem hiding this comment.
The spatial_axis indices for RandFlip are hardcoded assuming a channel-last layout ([D, H, W, C]). If the pipeline is configured with ensure_channel_last=False, the layout becomes [C, D, H, W], and these indices will incorrectly target the channel dimension or the wrong spatial dimensions. Consider determining the spatial axes dynamically based on the ensure_channel_last setting.
Comment on lines
+185
to
+192
| from scipy.ndimage import zoom as ndimage_zoom | ||
|
|
||
| ds_label = ndimage_zoom( | ||
| label_batch.astype(np.float32), | ||
| zoom_factors, | ||
| order=0, | ||
| mode="nearest", | ||
| ).astype(label_batch.dtype) |
There was a problem hiding this comment.
Importing scipy.ndimage.zoom inside __getitem__ and performing CPU-based resizing for every batch is inefficient and will significantly bottleneck the training pipeline, especially for 3D volumes. It is recommended to move the import to the top of the file and consider using faster resizing operations (like ops.image.resize for 2D or a custom op for 3D) or pre-calculating multi-scale targets during preprocessing.
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (5)
medicai/trainer/nnunet/pipeline.py (2)
45-63: Consider logging the exception instead of silently passing.The bare
except Exception: passat lines 61-62 silently swallows all errors, making debugging difficult. Consider logging the exception at debug level.Suggested improvement
except Exception: - pass + import logging + logging.debug("nvidia-smi query failed, using default GPU memory", exc_info=True) return float(default)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/pipeline.py` around lines 45 - 63, The _auto_detect_gpu_memory function currently swallows all exceptions with a bare except; change it to catch Exception as e and log the error at debug level (e.g., logger.debug("Failed to auto-detect GPU memory", exc_info=True)) so failures are visible during debugging; if no module logger exists, create one with logging.getLogger(__name__) at module scope and use that; keep the existing fallback behavior of returning float(default).
456-460: Consider using iterable unpacking for clarity.Static analysis suggests using
[1, *patch_size, n_mod]instead of list concatenation.Suggested fix
- dummy = np.zeros([1] + patch_size + [n_mod], dtype=np.float32) + dummy = np.zeros([1, *patch_size, n_mod], dtype=np.float32)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/pipeline.py` around lines 456 - 460, Replace the list concatenation used to build the dummy input shape with iterable unpacking for clarity: where the code constructs dummy = np.zeros([1] + patch_size + [n_mod], dtype=np.float32) (in the block that sets n_outputs from net_cfg and creates dummy before calling model), change it to use a single list literal with unpacking like [1, *patch_size, n_mod] so the dummy shape is clearer and avoids concatenating lists.medicai/trainer/nnunet/planning/planners.py (2)
1-19: Module docstring placement is unusual.The module docstring appears after the
from collections.abc import Sequenceimport (lines 3-16) rather than at the very beginning of the file. While this works, the conventional pattern is to have the docstring first.Suggested fix
+""" +nnunet_keras/planning/patch_size_planner.py +============================================ +Heuristic patch size computation following nnU-Net rules: +... +""" + from collections.abc import Sequence - -""" -nnunet_keras/planning/patch_size_planner.py -============================================ -... -""" import math🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/planning/planners.py` around lines 1 - 19, The module-level docstring is placed after the import of Sequence; move the triple-quoted nnunet_keras/planning/patch_size_planner.py docstring to the very top of the file (before any imports) so it becomes the module docstring, ensuring the docstring text remains intact and the import "from collections.abc import Sequence" and other imports follow it.
715-737: Stub implementations noted.The
nnUNetPlannerResEncMandnnUNetPlannerResEncLclasses are currently stubs with TODO comments. This is acceptable for the initial implementation but should be tracked for follow-up.Would you like me to open an issue to track implementing the ResEnc-specific architecture modifications?
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/planning/planners.py` around lines 715 - 737, The two planner classes nnUNetPlannerResEncM and nnUNetPlannerResEncL are left as silent stubs; instead of leaving TODO comments, update them so they explicitly signal unimplemented behavior and are tracked: either implement the ResEnc-specific modifications (e.g., adjust base_filters/max_filters/decoder heads inside __init__ of nnUNetPlannerResEncM and nnUNetPlannerResEncL) or, if implementation is deferred, raise a clear NotImplementedError in each __init__ and add a one-line TODO with the issue tracker ID (create an issue first if none exists) so future work is linked to that issue; ensure references to the class names (nnUNetPlannerResEncM, nnUNetPlannerResEncL) appear in the TODO/exception message for easy traceability.medicai/trainer/nnunet/data/dataset_fingerprint.py (1)
272-275: Addstrict=Trueto zip for consistency.While
np.uniqueguarantees matching lengths, addingstrict=Trueimproves code consistency with other zip usages in this PR.Suggested fix
- for val, count in zip(unique, counts): + for val, count in zip(unique, counts, strict=True):🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/data/dataset_fingerprint.py` around lines 272 - 275, The zip over outputs from np.unique in the loop using variables unique and counts should be called with strict=True for consistency; update the loop in dataset_fingerprint.py where unique, counts = np.unique(label_data.astype(np.int64), return_counts=True) and the subsequent for val, count in zip(unique, counts): to use zip(unique, counts, strict=True) so it enforces equal-length pairing when updating class_voxel_counts[int(val)].
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@medicai/trainer/nnunet/data/preprocessing.py`:
- Around line 594-618: The code checks manifest_file but never loads it and then
references undefined variables manifest and items; fix by loading the manifest
from manifest_file into a manifest variable (e.g., parse JSON/YAML into a dict)
before creating worker_fn, extract the list of items from that manifest (e.g.,
manifest["items"] or manifest.get("items")) and validate that items is a
non-empty iterable, and then use that items variable in the ProcessPoolExecutor
map and the serial loop; ensure any missing keys raise a clear ValueError.
Reference: manifest_file, manifest, items, _process_item_helper, worker_fn.
In `@medicai/trainer/nnunet/training/metrics.py`:
- Around line 33-37: Validate that y_true_np and y_pred_np have identical shapes
before computing masks for class_idx in the metric function (where y_true_np,
y_pred_np and class_idx are used); if shapes differ, raise a clear ValueError
(include the shapes in the message) so mismatched arrays don't silently produce
incorrect masks and distances, and only then proceed to compute mask_true,
mask_pred and the existing early-return of float("inf") when either mask sums to
zero.
In `@medicai/trainer/nnunet/utils/io.py`:
- Around line 275-286: The saving block that handles raster formats (the if
path_str.endswith... branch) currently casts int16/int32/int64 to np.uint8
directly (via image_out = image_out.astype(np.uint8)), causing silent data loss;
modify this branch in the io.py save routine so that before casting you either
(a) explicitly normalize/scale image_out to the 0-255 range (e.g., compute
min/max and scale to uint8) when values exceed [0,255], or (b) raise or log a
clear warning/error if automatic scaling is undesirable; update the logic around
the image_out/dtype handling (the image_out = image.astype(dtype) and subsequent
dtype check) to perform scaling or emit warnings prior to skio.imsave, ensuring
grayscale channel squeezing still occurs and that the final dtype is uint8 only
after safe conversion.
---
Nitpick comments:
In `@medicai/trainer/nnunet/data/dataset_fingerprint.py`:
- Around line 272-275: The zip over outputs from np.unique in the loop using
variables unique and counts should be called with strict=True for consistency;
update the loop in dataset_fingerprint.py where unique, counts =
np.unique(label_data.astype(np.int64), return_counts=True) and the subsequent
for val, count in zip(unique, counts): to use zip(unique, counts, strict=True)
so it enforces equal-length pairing when updating class_voxel_counts[int(val)].
In `@medicai/trainer/nnunet/pipeline.py`:
- Around line 45-63: The _auto_detect_gpu_memory function currently swallows all
exceptions with a bare except; change it to catch Exception as e and log the
error at debug level (e.g., logger.debug("Failed to auto-detect GPU memory",
exc_info=True)) so failures are visible during debugging; if no module logger
exists, create one with logging.getLogger(__name__) at module scope and use
that; keep the existing fallback behavior of returning float(default).
- Around line 456-460: Replace the list concatenation used to build the dummy
input shape with iterable unpacking for clarity: where the code constructs dummy
= np.zeros([1] + patch_size + [n_mod], dtype=np.float32) (in the block that sets
n_outputs from net_cfg and creates dummy before calling model), change it to use
a single list literal with unpacking like [1, *patch_size, n_mod] so the dummy
shape is clearer and avoids concatenating lists.
In `@medicai/trainer/nnunet/planning/planners.py`:
- Around line 1-19: The module-level docstring is placed after the import of
Sequence; move the triple-quoted nnunet_keras/planning/patch_size_planner.py
docstring to the very top of the file (before any imports) so it becomes the
module docstring, ensuring the docstring text remains intact and the import
"from collections.abc import Sequence" and other imports follow it.
- Around line 715-737: The two planner classes nnUNetPlannerResEncM and
nnUNetPlannerResEncL are left as silent stubs; instead of leaving TODO comments,
update them so they explicitly signal unimplemented behavior and are tracked:
either implement the ResEnc-specific modifications (e.g., adjust
base_filters/max_filters/decoder heads inside __init__ of nnUNetPlannerResEncM
and nnUNetPlannerResEncL) or, if implementation is deferred, raise a clear
NotImplementedError in each __init__ and add a one-line TODO with the issue
tracker ID (create an issue first if none exists) so future work is linked to
that issue; ensure references to the class names (nnUNetPlannerResEncM,
nnUNetPlannerResEncL) appear in the TODO/exception message for easy
traceability.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: c34726e5-c80a-428c-9a12-1704202a471d
📒 Files selected for processing (11)
medicai/models/nnunet/blocks.pymedicai/models/nnunet/dynamic_unet.pymedicai/models/nnunet/unet.pymedicai/trainer/nnunet/data/dataset_fingerprint.pymedicai/trainer/nnunet/data/manifest.pymedicai/trainer/nnunet/data/preprocessing.pymedicai/trainer/nnunet/pipeline.pymedicai/trainer/nnunet/planning/planners.pymedicai/trainer/nnunet/training/metrics.pymedicai/trainer/nnunet/utils/config.pymedicai/trainer/nnunet/utils/io.py
✅ Files skipped from review due to trivial changes (2)
- medicai/models/nnunet/dynamic_unet.py
- medicai/models/nnunet/blocks.py
Comment on lines
+33
to
+37
| mask_true = (y_true_np == class_idx).astype(np.uint8) | ||
| mask_pred = (y_pred_np == class_idx).astype(np.uint8) | ||
|
|
||
| if mask_true.sum() == 0 or mask_pred.sum() == 0: | ||
| return float("inf") |
There was a problem hiding this comment.
Add shape validation to prevent silent mismatches.
If y_true_np and y_pred_np have different shapes, the comparison y_true_np == class_idx and y_pred_np == class_idx will succeed but produce masks of different shapes, leading to incorrect distance calculations without any error.
🛡️ Proposed fix
if not _SCIPY_AVAILABLE:
return float("nan")
+ if y_true_np.shape != y_pred_np.shape:
+ raise ValueError(
+ f"Shape mismatch: y_true_np {y_true_np.shape} vs y_pred_np {y_pred_np.shape}"
+ )
+
mask_true = (y_true_np == class_idx).astype(np.uint8)
mask_pred = (y_pred_np == class_idx).astype(np.uint8)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/trainer/nnunet/training/metrics.py` around lines 33 - 37, Validate
that y_true_np and y_pred_np have identical shapes before computing masks for
class_idx in the metric function (where y_true_np, y_pred_np and class_idx are
used); if shapes differ, raise a clear ValueError (include the shapes in the
message) so mismatched arrays don't silently produce incorrect masks and
distances, and only then proceed to compute mask_true, mask_pred and the
existing early-return of float("inf") when either mask sums to zero.
Comment on lines
+275
to
+286
| if path_str.endswith((".png", ".jpg", ".jpeg", ".tif", ".tiff")): | ||
| path = Path(path) | ||
| path.parent.mkdir(parents=True, exist_ok=True) | ||
| image_out = image.astype(dtype) | ||
| if image_out.ndim == 3 and image_out.shape[0] <= 4: | ||
| image_out = image_out.transpose(1, 2, 0) | ||
| if dtype in (np.int16, np.int32, np.int64) and path_str.endswith((".png", ".jpg", ".jpeg")): | ||
| image_out = image_out.astype(np.uint8) | ||
| if image_out.ndim == 3 and image_out.shape[-1] == 1: | ||
| image_out = np.squeeze(image_out, axis=-1) | ||
| skio.imsave(str(path), image_out, check_contrast=False) | ||
| return |
There was a problem hiding this comment.
Potential data loss during raster image saving.
The dtype conversion at lines 281-282 casts int16/int32/int64 directly to uint8 without scaling. This will cause data loss for values outside [0, 255]. Consider adding normalization or warning the user.
Suggested improvement
if dtype in (np.int16, np.int32, np.int64) and path_str.endswith((".png", ".jpg", ".jpeg")):
- image_out = image_out.astype(np.uint8)
+ # Clip and convert - warn if values are outside uint8 range
+ if image_out.min() < 0 or image_out.max() > 255:
+ import warnings
+ warnings.warn(
+ f"Image values [{image_out.min()}, {image_out.max()}] will be clipped to [0, 255] for {path}",
+ UserWarning
+ )
+ image_out = np.clip(image_out, 0, 255).astype(np.uint8)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/trainer/nnunet/utils/io.py` around lines 275 - 286, The saving block
that handles raster formats (the if path_str.endswith... branch) currently casts
int16/int32/int64 to np.uint8 directly (via image_out =
image_out.astype(np.uint8)), causing silent data loss; modify this branch in the
io.py save routine so that before casting you either (a) explicitly
normalize/scale image_out to the 0-255 range (e.g., compute min/max and scale to
uint8) when values exceed [0,255], or (b) raise or log a clear warning/error if
automatic scaling is undesirable; update the logic around the image_out/dtype
handling (the image_out = image.astype(dtype) and subsequent dtype check) to
perform scaling or emit warnings prior to skio.imsave, ensuring grayscale
channel squeezing still occurs and that the final dtype is uint8 only after safe
conversion.
Owner
Author
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a Keras 3 implementation of the nnU-Net workflow, featuring automated dataset fingerprinting, heuristic planning, and a dynamic U-Net architecture. Key feedback identifies critical bugs, including uninitialized variables in the fingerprinting module and the use of 2D resizing operations on 3D volumetric data during deep supervision. Further issues were noted regarding potential dimension mismatches in multi-label resampling, performance bottlenecks from CPU-based interpolation, and inconsistent EMA configurations in the trainer that could cause runtime errors.
Comment on lines
+153
to
+155
| spacings = [] | ||
| sizes = [] | ||
| median_relative_sizes = [] |
There was a problem hiding this comment.
The variables spatial_dims, total_voxels, class_voxel_counts, and all_images_per_modality are used throughout the fingerprint_dataset function but are never initialized, which will lead to a NameError or UnboundLocalError upon execution.
spacings = []
sizes = []
median_relative_sizes = []
spatial_dims = None
total_voxels = 0
class_voxel_counts = {i: 0 for i in range(n_classes)}
all_images_per_modality = {i: [] for i in range(len(modalities))}
Comment on lines
+173
to
+179
| interp = "bilinear" if self.spatial_dims == 2 else "trilinear" | ||
| resized = ops.image.resize( | ||
| seg_outputs[i], | ||
| size=target_shape, | ||
| interpolation=interp, | ||
| ) | ||
| out_dict[f"aux_{i-1}"] = resized |
There was a problem hiding this comment.
ops.image.resize is specifically designed for 2D images (4D tensors). In a 3D segmentation pipeline, seg_outputs are 5D tensors (B, D, H, W, C), and this operation will fail or produce incorrect results. Consider using a 3D-capable resizing utility or a custom layer that handles volumetric data.
| try: | ||
| with open(prop_file, "r", encoding="utf-8") as f: | ||
| self.properties_map[str(cf)] = json.load(f) | ||
| except Exception: |
There was a problem hiding this comment.
Avoid catching generic Exception and passing silently. This can mask critical issues such as file access errors, malformed JSON in property files, or unexpected data structures, making debugging difficult.
Comment on lines
+190
to
+196
| ds_label = ndimage_zoom( | ||
| label_batch.astype(np.float32), | ||
| zoom_factors, | ||
| order=0, | ||
| mode="nearest", | ||
| ).astype(label_batch.dtype) | ||
| y_dict[f"aux_{i}"] = ds_label |
There was a problem hiding this comment.
ndimage_zoom requires the length of zoom_factors to match the number of dimensions in the input array. For multi-label tasks, label_batch is a 5D tensor (B, D, H, W, C), but zoom_factors only contains 4 elements, which will cause a runtime error. Additionally, performing CPU-based interpolation on large 3D volumes during every batch fetch can significantly bottleneck training throughput.
Comment on lines
+113
to
+125
| if isinstance(cfg.optimizer, str): | ||
| if cfg.optimizer == "sgd": | ||
| return keras.optimizers.SGD( | ||
| learning_rate=lr, | ||
| momentum=cfg.momentum, | ||
| nesterov=cfg.nesterov, | ||
| weight_decay=cfg.weight_decay, | ||
| gradient_accumulation_steps=cfg.gradient_accumulation_steps, | ||
| global_clipnorm=12.0, | ||
| use_ema=cfg.use_ema, | ||
| ema_momentum=cfg.ema_momentum, | ||
| ) | ||
| return cfg.optimizer |
There was a problem hiding this comment.
The EMA configuration is inconsistent. While use_ema is a general configuration parameter, it is only applied during optimizer construction if the optimizer is explicitly set to the string "sgd". If a user specifies another optimizer (e.g., "adam"), the SwapEMAWeights callback is still added at line 257, which will cause a crash because the optimizer will lack the necessary EMA attributes.
There was a problem hiding this comment.
Actionable comments posted: 12
🧹 Nitpick comments (10)
medicai/dataloader/nnunet/augmentations.py (1)
13-25: Most ofAugmentationConfigis unused right now.
p_scale,p_elastic,p_gamma,p_noise,scale_range,elastic_*,gamma_range,noise_variance, andmirror_axesare declared but never consulted; the pipeline always builds the same hard-coded flip/rotate chain. Either wire these fields into the transform sequence or drop them until they’re actually supported.Also applies to: 48-54
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/dataloader/nnunet/augmentations.py` around lines 13 - 25, AugmentationConfig contains many unused fields (p_scale, p_elastic, p_gamma, p_noise, scale_range, elastic_alpha, elastic_sigma, gamma_range, noise_variance, mirror_axes) but the transform pipeline always uses a hard-coded flip/rotate chain; update the transform builder (the function that constructs the augmentation pipeline—look for where AugmentationConfig is consumed, e.g., the pipeline or get_training_transforms/build_transforms function) to conditionally add/parameterize scale, elastic, gamma, noise, and mirror transforms using their probability flags and range/hyperparameter fields, wiring p_* to each transform's application probability and using scale_range, elastic_alpha/elastic_sigma, gamma_range, noise_variance, and mirror_axes for parameters; alternatively remove these unused fields from AugmentationConfig if you choose not to support them yet.medicai/dataloader/nnunet/preprocessing.py (1)
121-131: Consider using spread syntax for tuple construction.The pattern
padded_bbox + (slice(None),)can be more clearly written as(*padded_bbox, slice(None)).♻️ Proposed fix
if image.ndim == len(shape): cropped_image = image[padded_bbox] else: - cropped_image = image[padded_bbox + (slice(None),)] + cropped_image = image[(*padded_bbox, slice(None))] if label is None: cropped_label = None elif label.ndim == len(shape): cropped_label = label[padded_bbox] else: - cropped_label = label[padded_bbox + (slice(None),)] + cropped_label = label[(*padded_bbox, slice(None))]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/dataloader/nnunet/preprocessing.py` around lines 121 - 131, Replace the tuple-concatenation pattern used for channel-aware indexing with tuple unpacking for clarity: wherever the code constructs padded_bbox + (slice(None),) (used when selecting from image and label in preprocessing.py, producing cropped_image and cropped_label), change it to use the spread/unpack form (*padded_bbox, slice(None)) so the indexing intent is clearer and consistent for both image and label branches.medicai/dataloader/nnunet/manifest.py (1)
131-132: Clarify theor Nonepattern for empty regions.The expression
[[int(v) for v in region] for region in item_dict.get("regions", [])] or Noneconverts an empty list toNone. This is intentional but could be made more explicit for readability.♻️ Proposed fix for clarity
- regions=[[int(v) for v in region] for region in item_dict.get("regions", [])] - or None, + regions=( + [[int(v) for v in region] for region in item_dict.get("regions", [])] + or None + ),🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/dataloader/nnunet/manifest.py` around lines 131 - 132, The current expression regions=[[int(v) for v in region] for region in item_dict.get("regions", [])] or None implicitly converts an empty list to None which is unclear; make this explicit by first retrieving the raw regions (e.g., raw_regions = item_dict.get("regions", [])), build the converted list with [[int(v) for v in region] for region in raw_regions], and then set regions to None if that converted list is empty (e.g., regions = converted if converted else None) so the intent is readable and unambiguous; update the assignment in manifest.py where regions is constructed to follow this pattern and reference the same variable names used there.medicai/dataloader/nnunet/dataset_fingerprint.py (1)
244-247: Addstrict=Trueto zip call for safety.The
zip(unique, counts)should usestrict=Truesince both arrays come fromnp.unique(..., return_counts=True)and must have matching lengths.♻️ Proposed fix
- for val, count in zip(unique, counts): + for val, count in zip(unique, counts, strict=True):🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/dataloader/nnunet/dataset_fingerprint.py` around lines 244 - 247, The loop iterating over unique and counts should use zip(..., strict=True) to ensure lengths match; update the zip(unique, counts) call in dataset_fingerprint.py (the block that updates class_voxel_counts and total_voxels using unique, counts, class_voxel_counts, and cropped_shape) to zip(unique, counts, strict=True) so mismatched lengths raise immediately and guard against silent bugs.medicai/trainer/nnunet/planning/planners.py (3)
199-199: Remove duplicateimport mathstatement.
mathis already imported at line 1. This duplicate import should be removed.♻️ Proposed fix
-import math - # Constants (tunable)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/planning/planners.py` at line 199, Duplicate import of the math module was added (the second "import math"); remove the redundant "import math" statement so only the original import at the top of planners.py remains, ensuring no other code references change—look for the duplicate "import math" line to delete.
644-665: TODO: ResEncM and ResEncL planners are stubs.Both
nnUNetPlannerResEncMandnnUNetPlannerResEncLclasses are currently identical to the base planner with only TODO comments. These should either be implemented or removed if not needed for this PR.Would you like me to help implement the ResEnc-specific modifications (e.g., adjusted
base_filters,max_filters) or open an issue to track this?🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/planning/planners.py` around lines 644 - 665, The two planner classes nnUNetPlannerResEncM and nnUNetPlannerResEncL are placeholders and should either be implemented with their ResEnc-specific parameters or removed; update nnUNetPlannerResEncM to override the base __init__ to set ResEncM defaults (e.g., increase base_filters, adjust max_filters, and any architecture flags used by nnUNetPlanner such as encoder_type or residual_blocks) and update nnUNetPlannerResEncL similarly with larger base_filters/more decoder heads, ensuring you call super().__init__(...) then assign/override the planner attributes (e.g., base_filters, max_filters, decoder_heads, encoder_type) used by the rest of the planner logic; if these variants are not needed for this PR, remove both classes to avoid dead/stub code and open an issue to track adding ResEnc variants later.
312-317: Move imports to top of file.These imports (
Path,numpy, and local config imports) should be consolidated with the other imports at the top of the file per PEP 8 conventions.♻️ Proposed fix
Move these lines to the top of the file after
import math:from pathlib import Path import numpy as np from medicai.dataloader.nnunet.resampling import compute_zoom_factors from medicai.trainer.nnunet.utils.config import DatasetFingerprint, NetworkConfig, nnUNetPlan🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/planning/planners.py` around lines 312 - 317, The listed imports (Path, numpy, compute_zoom_factors, DatasetFingerprint, NetworkConfig, nnUNetPlan) are located mid-file; move them to the module import section so they appear with the other top-level imports (per PEP8). Specifically, relocate the four lines referencing Path, numpy, medicai.dataloader.nnunet.resampling.compute_zoom_factors, and medicai.trainer.nnunet.utils.config.{DatasetFingerprint, NetworkConfig, nnUNetPlan} into the top of the file’s import block (place them immediately after the existing import math) and remove the duplicate mid-file import statements.medicai/trainer/nnunet/pipeline.py (1)
319-321: Use spread syntax for cleaner list construction.The static analysis suggestion to use
[1, *patch_size, n_mod]is valid and more readable than concatenation.♻️ Proposed fix
- dummy = np.zeros([1] + patch_size + [n_mod], dtype=np.float32) + dummy = np.zeros([1, *patch_size, n_mod], dtype=np.float32)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/trainer/nnunet/pipeline.py` around lines 319 - 321, The list used to construct the dummy input shape is built by concatenation (np.zeros([1] + patch_size + [n_mod], ...)); replace it with Python's spread/unpacking syntax for clarity by constructing the shape as [1, *patch_size, n_mod] when creating dummy in the block that calls model(dummy, training=False) and then model.load_weights(...); update the dummy creation only so patch_size and n_mod are unchanged and behavior remains the same.medicai/dataloader/nnunet/dataset.py (2)
119-126: Consider addingstrict=Trueto zip calls for defensive coding.While the lists
pad_beforeandpad_afterare built in the same loop and should always have matching lengths, addingstrict=Truewould catch any future bugs if the construction logic changes.♻️ Proposed fix
- pad_width_img = [(pb, pa) for pb, pa in zip(pad_before, pad_after)] + [(0, 0)] + pad_width_img = [(pb, pa) for pb, pa in zip(pad_before, pad_after, strict=True)] + [(0, 0)]- pad_width_lbl = [(pb, pa) for pb, pa in zip(pad_before, pad_after)] + pad_width_lbl = [(pb, pa) for pb, pa in zip(pad_before, pad_after, strict=True)]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/dataloader/nnunet/dataset.py` around lines 119 - 126, The zip calls building pad_width_img and pad_width_lbl should be made defensive by adding strict=True to both zip(...) invocations (i.e., use zip(pad_before, pad_after, strict=True)) so mismatched pad_before/pad_after lengths raise immediately; update the two occurrences that construct pad_width_img and pad_width_lbl (and any other zip(pad_before, pad_after) in this scope) to include strict=True and run tests on Python 3.10+ to ensure compatibility.
44-48: Consider logging failed property file loads for debugging.Silently swallowing exceptions when loading property files can make debugging difficult if there are subtle file format issues or permission problems.
♻️ Proposed fix
+import logging + +logger = logging.getLogger(__name__) + ... if prop_file.exists(): try: with open(prop_file, "r", encoding="utf-8") as f: self.properties_map[str(cf)] = json.load(f) - except Exception: - pass + except Exception as e: + logger.debug("Failed to load properties for %s: %s", cf, e)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@medicai/dataloader/nnunet/dataset.py` around lines 44 - 48, Replace the silent except in the property-file load loop so that exceptions are caught as `except Exception as e` and the failure is logged with the property filename and the exception details (including stacktrace) instead of being silently ignored; update the code around `self.properties_map`, `prop_file`, and `cf` to call your module/class logger (e.g. `self.logger` or `logging.getLogger(__name__)`) to emit a clear error message mentioning `prop_file`/`cf` and `e` so you can debug file-format or permission issues.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@medicai/dataloader/nnunet/cross_validation.py`:
- Around line 138-146: The aggregation loop currently uses enumerate(values)
which reindexes folds and can mislabel per-fold metrics; instead iterate over
the original fold_results so the fold index matches the source fold: for each
key in sorted(all_keys) build values as before but replace the inner loop with
"for fold_idx, d in enumerate(fold_results): if key in d: v = d[key];
summary[f'{key}_fold_{fold_idx}'] = float(v)" (keep summary, all_keys,
fold_results and arr.mean()/arr.std() logic unchanged).
- Around line 8-18: The normalize_case_id function is erroneously removing any
trailing _dddd numeric suffix (re.sub(r"_\d{4}$", "", ...)) which collapses
distinct preprocessed stems like patient_2024 into patient; update
normalize_case_id (used by _build_datasets and split matching) to stop stripping
the trailing four-digit pattern—return the normalized stem after removing file
extensions only and do not apply the re.sub that removes _\d{4}$ so that
preprocessed case IDs remain distinct.
In `@medicai/dataloader/nnunet/dataset_fingerprint.py`:
- Around line 153-156: In fingerprint_dataset initialize the missing variables
before the main fingerprinting loop: set spatial_dims to None or 0 as a
placeholder, set class_voxel_counts to a zeros container (e.g., list of zeros or
defaultdict(int)) sized for num_classes, set total_voxels to 0, and set
all_images_per_modality to an empty list or dict (e.g., list of lists or
defaultdict(list)) depending on how images are aggregated; update subsequent
code to use these containers (functions/variables referenced:
fingerprint_dataset, spatial_dims, class_voxel_counts, total_voxels,
all_images_per_modality) so no NameError occurs at runtime.
In `@medicai/dataloader/nnunet/normalization.py`:
- Around line 161-172: compute_intensity_stats currently accumulates every
flattened image in voxels_list and then concatenates them, which duplicates
memory usage; replace this with a streaming/stateless approach: compute mean and
variance online (e.g., Welford's algorithm) over images in the loop to produce
mean and std without storing full arrays, and concurrently perform reservoir
sampling (or fixed-size random subsampling per image) to collect a bounded
sample set used only for percentile computations; update references to
nonzero_only and images handling inside compute_intensity_stats so you never
call np.concatenate on all voxels and only compute percentiles from the sampled
array while returning mean/std from the streaming aggregator.
- Around line 84-104: The current flow computes mean/std from foreground (fg)
but applies the affine transform to the whole image; change the logic in the
normalization block (referencing self.nonzero_only, mask,
self.use_mask_for_norm, fg, mean, std, image) so that when foreground-only stats
are used you only normalize the foreground voxels and leave background voxels as
zero. Concretely, compute mean/std from fg as before, then compute normalized_fg
= (fg - mean) / std and write those values back into the corresponding positions
of image (using mask > 0 or image != 0 to find indices) and return image; when
not using nonzero_only keep the existing whole-image normalization path.
In `@medicai/models/nnunet/unet.py`:
- Around line 166-182: The method currently gates returning the deep-supervision
dict on the local `training` flag, causing the model output type to change
between train and eval; instead, always return a dict when
`self.deep_supervision` is True. Modify the branch in the function that builds
`out_dict` (using `seg_outputs`, `ops.shape`, `ops.image.resize`, and `interp`)
so that it executes whenever `self.deep_supervision` is enabled (regardless of
`training`), and only return `seg_outputs[0]` for the non-deep-supervision case.
- Around line 169-178: The deep-supervision resize uses ops.image.resize with
"trilinear" which fails for 3D; update the loop in unet (where seg_outputs and
self.spatial_dims are used) to call medicai.utils.image.resize_volumes (or
medicai.transforms.resize.resize_volumes) when self.spatial_dims != 2, passing
the same target_shape and using the trilinear method/appropriate argument names,
and keep ops.image.resize for the 2D (bilinear) case; also add the necessary
import for resize_volumes at the top of the module.
In `@medicai/trainer/nnunet/pipeline.py`:
- Around line 1-41: The module is missing the numpy import used by predict and
_postprocess_prediction, causing NameError; add an import for numpy (e.g.,
import numpy as np) at the top of the file alongside the other imports so
references to np in predict(...) and _postprocess_prediction(...) resolve
correctly.
In `@medicai/trainer/nnunet/training/trainer.py`:
- Around line 262-268: The training run currently always validates every epoch
because self.model.fit is called with validation_data=self.val_dataset and
doesn't respect cfg.val_every_n_epochs; update run() to pass validation_freq (or
implement an equivalent callback) so validation only runs every
cfg.val_every_n_epochs epochs when cfg.val_every_n_epochs > 1, referencing
self.model.fit, cfg.val_every_n_epochs, self.train_dataset and self.val_dataset
to conditionally set validation_freq (or wrap/replace validation with a callback
that runs validation every N epochs) and leave behavior unchanged when
cfg.val_every_n_epochs == 1 or None.
- Around line 80-91: The ModelCheckpoint setup is using save_freq as epochs and
forcing mode="max", which is incorrect; in the checkpoint creation (where
ModelCheckpoint is instantiated) multiply cfg.save_every_n_epochs by
cfg.iters_per_epoch so save_freq = cfg.save_every_n_epochs * cfg.iters_per_epoch
(handle None/0 safely), and change mode="max" to mode="auto" so Keras infers
min/max from the metric name produced by _metric_monitor_name (which returns
"val_loss"/"val_final_loss" when no metrics are present). Ensure references to
cfg.save_every_n_epochs, cfg.iters_per_epoch, ModelCheckpoint and
_metric_monitor_name are updated accordingly.
In `@medicai/trainer/nnunet/utils/config.py`:
- Around line 120-141: The defaults for 2D/3D are hardcoded to 3D sizes; update
the constructor of the NetworkConfig-like class to choose defaults based on
spatial_dims: set patch_size to [128]*spatial_dims when patch_size is None,
kernel_size to [3]*spatial_dims when kernel_size is None, and set
pool_op_kernel_sizes to a sensible default matching spatial_dims (e.g., empty
list or list of kernel tuples sized by spatial_dims) when pool_op_kernel_sizes
is None; modify the initialization logic around the attributes patch_size,
kernel_size, and pool_op_kernel_sizes in the __init__ (the block referencing
self.patch_size, self.kernel_size, self.pool_op_kernel_sizes) so callers with
spatial_dims=2 get 2D shapes and spatial_dims=3 get 3D shapes.
In `@medicai/trainer/nnunet/utils/io.py`:
- Around line 48-49: The 4D write path in utils/io.py uses
image.transpose(3,2,1,0) which is not the inverse of load_nifti(); change the
transpose axes in the image.ndim == 4 branch so that save mirrors load_nifti()'s
(2,1,0,3) mapping (i.e., reverse the same axes order used in load_nifti()) to
preserve multi-channel NIfTI layout for the function that handles saving (the
branch handling image.ndim == 4).
---
Nitpick comments:
In `@medicai/dataloader/nnunet/augmentations.py`:
- Around line 13-25: AugmentationConfig contains many unused fields (p_scale,
p_elastic, p_gamma, p_noise, scale_range, elastic_alpha, elastic_sigma,
gamma_range, noise_variance, mirror_axes) but the transform pipeline always uses
a hard-coded flip/rotate chain; update the transform builder (the function that
constructs the augmentation pipeline—look for where AugmentationConfig is
consumed, e.g., the pipeline or get_training_transforms/build_transforms
function) to conditionally add/parameterize scale, elastic, gamma, noise, and
mirror transforms using their probability flags and range/hyperparameter fields,
wiring p_* to each transform's application probability and using scale_range,
elastic_alpha/elastic_sigma, gamma_range, noise_variance, and mirror_axes for
parameters; alternatively remove these unused fields from AugmentationConfig if
you choose not to support them yet.
In `@medicai/dataloader/nnunet/dataset_fingerprint.py`:
- Around line 244-247: The loop iterating over unique and counts should use
zip(..., strict=True) to ensure lengths match; update the zip(unique, counts)
call in dataset_fingerprint.py (the block that updates class_voxel_counts and
total_voxels using unique, counts, class_voxel_counts, and cropped_shape) to
zip(unique, counts, strict=True) so mismatched lengths raise immediately and
guard against silent bugs.
In `@medicai/dataloader/nnunet/dataset.py`:
- Around line 119-126: The zip calls building pad_width_img and pad_width_lbl
should be made defensive by adding strict=True to both zip(...) invocations
(i.e., use zip(pad_before, pad_after, strict=True)) so mismatched
pad_before/pad_after lengths raise immediately; update the two occurrences that
construct pad_width_img and pad_width_lbl (and any other zip(pad_before,
pad_after) in this scope) to include strict=True and run tests on Python 3.10+
to ensure compatibility.
- Around line 44-48: Replace the silent except in the property-file load loop so
that exceptions are caught as `except Exception as e` and the failure is logged
with the property filename and the exception details (including stacktrace)
instead of being silently ignored; update the code around `self.properties_map`,
`prop_file`, and `cf` to call your module/class logger (e.g. `self.logger` or
`logging.getLogger(__name__)`) to emit a clear error message mentioning
`prop_file`/`cf` and `e` so you can debug file-format or permission issues.
In `@medicai/dataloader/nnunet/manifest.py`:
- Around line 131-132: The current expression regions=[[int(v) for v in region]
for region in item_dict.get("regions", [])] or None implicitly converts an empty
list to None which is unclear; make this explicit by first retrieving the raw
regions (e.g., raw_regions = item_dict.get("regions", [])), build the converted
list with [[int(v) for v in region] for region in raw_regions], and then set
regions to None if that converted list is empty (e.g., regions = converted if
converted else None) so the intent is readable and unambiguous; update the
assignment in manifest.py where regions is constructed to follow this pattern
and reference the same variable names used there.
In `@medicai/dataloader/nnunet/preprocessing.py`:
- Around line 121-131: Replace the tuple-concatenation pattern used for
channel-aware indexing with tuple unpacking for clarity: wherever the code
constructs padded_bbox + (slice(None),) (used when selecting from image and
label in preprocessing.py, producing cropped_image and cropped_label), change it
to use the spread/unpack form (*padded_bbox, slice(None)) so the indexing intent
is clearer and consistent for both image and label branches.
In `@medicai/trainer/nnunet/pipeline.py`:
- Around line 319-321: The list used to construct the dummy input shape is built
by concatenation (np.zeros([1] + patch_size + [n_mod], ...)); replace it with
Python's spread/unpacking syntax for clarity by constructing the shape as [1,
*patch_size, n_mod] when creating dummy in the block that calls model(dummy,
training=False) and then model.load_weights(...); update the dummy creation only
so patch_size and n_mod are unchanged and behavior remains the same.
In `@medicai/trainer/nnunet/planning/planners.py`:
- Line 199: Duplicate import of the math module was added (the second "import
math"); remove the redundant "import math" statement so only the original import
at the top of planners.py remains, ensuring no other code references change—look
for the duplicate "import math" line to delete.
- Around line 644-665: The two planner classes nnUNetPlannerResEncM and
nnUNetPlannerResEncL are placeholders and should either be implemented with
their ResEnc-specific parameters or removed; update nnUNetPlannerResEncM to
override the base __init__ to set ResEncM defaults (e.g., increase base_filters,
adjust max_filters, and any architecture flags used by nnUNetPlanner such as
encoder_type or residual_blocks) and update nnUNetPlannerResEncL similarly with
larger base_filters/more decoder heads, ensuring you call super().__init__(...)
then assign/override the planner attributes (e.g., base_filters, max_filters,
decoder_heads, encoder_type) used by the rest of the planner logic; if these
variants are not needed for this PR, remove both classes to avoid dead/stub code
and open an issue to track adding ResEnc variants later.
- Around line 312-317: The listed imports (Path, numpy, compute_zoom_factors,
DatasetFingerprint, NetworkConfig, nnUNetPlan) are located mid-file; move them
to the module import section so they appear with the other top-level imports
(per PEP8). Specifically, relocate the four lines referencing Path, numpy,
medicai.dataloader.nnunet.resampling.compute_zoom_factors, and
medicai.trainer.nnunet.utils.config.{DatasetFingerprint, NetworkConfig,
nnUNetPlan} into the top of the file’s import block (place them immediately
after the existing import math) and remove the duplicate mid-file import
statements.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b6dfe5ed-64da-4229-88ef-ba4240d2e2ef
📒 Files selected for processing (19)
.gitignoremedicai/dataloader/nnunet/__init__.pymedicai/dataloader/nnunet/augmentations.pymedicai/dataloader/nnunet/cross_validation.pymedicai/dataloader/nnunet/dataset.pymedicai/dataloader/nnunet/dataset_fingerprint.pymedicai/dataloader/nnunet/manifest.pymedicai/dataloader/nnunet/normalization.pymedicai/dataloader/nnunet/preprocessing.pymedicai/dataloader/nnunet/resampling.pymedicai/models/nnunet/blocks.pymedicai/models/nnunet/dynamic_unet.pymedicai/models/nnunet/unet.pymedicai/trainer/nnunet/__init__.pymedicai/trainer/nnunet/pipeline.pymedicai/trainer/nnunet/planning/planners.pymedicai/trainer/nnunet/training/trainer.pymedicai/trainer/nnunet/utils/config.pymedicai/trainer/nnunet/utils/io.py
✅ Files skipped from review due to trivial changes (3)
- .gitignore
- medicai/trainer/nnunet/init.py
- medicai/models/nnunet/dynamic_unet.py
🚧 Files skipped from review as they are similar to previous changes (1)
- medicai/models/nnunet/blocks.py
Comment on lines
+8
to
+18
| def normalize_case_id(case_id): | ||
| """Normalize identifiers by taking the stem of the file path. | ||
| The manifest handles all modality mappings, so we no longer need to | ||
| manually strip modality suffixes. | ||
| """ | ||
| normalized = Path(case_id).name | ||
| for suffix in (".nii.gz", ".nii", ".npz", ".png", ".jpg", ".jpeg", ".tif", ".tiff", ".dcm"): | ||
| if normalized.lower().endswith(suffix): | ||
| normalized = normalized[: -len(suffix)] | ||
| break | ||
| return re.sub(r"_\d{4}$", "", normalized) |
There was a problem hiding this comment.
Stop stripping every trailing _dddd from split IDs.
_build_datasets() normalizes already-preprocessed .npz stems before generating and matching splits, so a real case like patient_2024.npz becomes patient. That can merge distinct cases into one fold and later drop files when train/val membership is reconstructed.
🛠️ Proposed fix
- return re.sub(r"_\d{4}$", "", normalized)
+ return normalized📝 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
| def normalize_case_id(case_id): | |
| """Normalize identifiers by taking the stem of the file path. | |
| The manifest handles all modality mappings, so we no longer need to | |
| manually strip modality suffixes. | |
| """ | |
| normalized = Path(case_id).name | |
| for suffix in (".nii.gz", ".nii", ".npz", ".png", ".jpg", ".jpeg", ".tif", ".tiff", ".dcm"): | |
| if normalized.lower().endswith(suffix): | |
| normalized = normalized[: -len(suffix)] | |
| break | |
| return re.sub(r"_\d{4}$", "", normalized) | |
| def normalize_case_id(case_id): | |
| """Normalize identifiers by taking the stem of the file path. | |
| The manifest handles all modality mappings, so we no longer need to | |
| manually strip modality suffixes. | |
| """ | |
| normalized = Path(case_id).name | |
| for suffix in (".nii.gz", ".nii", ".npz", ".png", ".jpg", ".jpeg", ".tif", ".tiff", ".dcm"): | |
| if normalized.lower().endswith(suffix): | |
| normalized = normalized[: -len(suffix)] | |
| break | |
| return normalized |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/dataloader/nnunet/cross_validation.py` around lines 8 - 18, The
normalize_case_id function is erroneously removing any trailing _dddd numeric
suffix (re.sub(r"_\d{4}$", "", ...)) which collapses distinct preprocessed stems
like patient_2024 into patient; update normalize_case_id (used by
_build_datasets and split matching) to stop stripping the trailing four-digit
pattern—return the normalized stem after removing file extensions only and do
not apply the re.sub that removes _\d{4}$ so that preprocessed case IDs remain
distinct.
Comment on lines
+138
to
+146
| for key in sorted(all_keys): | ||
| values = [d[key] for d in fold_results if key in d] | ||
| if not values: | ||
| continue | ||
| arr = np.array(values, dtype=np.float64) | ||
| summary[f"mean_{key}"] = float(arr.mean()) | ||
| summary[f"std_{key}"] = float(arr.std()) | ||
| for fold_idx, v in enumerate(values): | ||
| summary[f"{key}_fold_{fold_idx}"] = float(v) |
There was a problem hiding this comment.
Keep the original fold numbers in aggregated metrics.
enumerate(values) renumbers only the folds that happened to emit key. If one fold is missing a metric, *_fold_1 can end up referring to the original fold 2, which makes downstream summaries misleading.
🛠️ Proposed fix
arr = np.array(values, dtype=np.float64)
summary[f"mean_{key}"] = float(arr.mean())
summary[f"std_{key}"] = float(arr.std())
- for fold_idx, v in enumerate(values):
- summary[f"{key}_fold_{fold_idx}"] = float(v)
+ for fold_idx, result in enumerate(fold_results):
+ if key in result:
+ summary[f"{key}_fold_{fold_idx}"] = float(result[key])📝 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
| for key in sorted(all_keys): | |
| values = [d[key] for d in fold_results if key in d] | |
| if not values: | |
| continue | |
| arr = np.array(values, dtype=np.float64) | |
| summary[f"mean_{key}"] = float(arr.mean()) | |
| summary[f"std_{key}"] = float(arr.std()) | |
| for fold_idx, v in enumerate(values): | |
| summary[f"{key}_fold_{fold_idx}"] = float(v) | |
| for key in sorted(all_keys): | |
| values = [d[key] for d in fold_results if key in d] | |
| if not values: | |
| continue | |
| arr = np.array(values, dtype=np.float64) | |
| summary[f"mean_{key}"] = float(arr.mean()) | |
| summary[f"std_{key}"] = float(arr.std()) | |
| for fold_idx, result in enumerate(fold_results): | |
| if key in result: | |
| summary[f"{key}_fold_{fold_idx}"] = float(result[key]) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/dataloader/nnunet/cross_validation.py` around lines 138 - 146, The
aggregation loop currently uses enumerate(values) which reindexes folds and can
mislabel per-fold metrics; instead iterate over the original fold_results so the
fold index matches the source fold: for each key in sorted(all_keys) build
values as before but replace the inner loop with "for fold_idx, d in
enumerate(fold_results): if key in d: v = d[key];
summary[f'{key}_fold_{fold_idx}'] = float(v)" (keep summary, all_keys,
fold_results and arr.mean()/arr.std() logic unchanged).
Comment on lines
+153
to
+156
| spacings = [] | ||
| sizes = [] | ||
| median_relative_sizes = [] | ||
|
|
There was a problem hiding this comment.
Critical: Multiple undefined variables will cause NameError at runtime.
The fingerprint_dataset function uses several variables that are never initialized:
spatial_dims(used at line 197 before first assignment)class_voxel_counts(used at lines 225, 245-246, 318)total_voxels(used at lines 226, 247)all_images_per_modality(used at lines 265, 285, 293)
These need to be initialized before the fingerprinting loop.
🐛 Proposed fix
spacings = []
sizes = []
median_relative_sizes = []
+ spatial_dims = None
+ class_voxel_counts = {i: 0 for i in range(len(class_names))}
+ total_voxels = 0
+ all_images_per_modality = {i: [] for i in range(len(modalities))}
for item in tqdm(items, desc="Fingerprinting"):📝 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
| spacings = [] | |
| sizes = [] | |
| median_relative_sizes = [] | |
| spacings = [] | |
| sizes = [] | |
| median_relative_sizes = [] | |
| spatial_dims = None | |
| class_voxel_counts = {i: 0 for i in range(len(class_names))} | |
| total_voxels = 0 | |
| all_images_per_modality = {i: [] for i in range(len(modalities))} |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/dataloader/nnunet/dataset_fingerprint.py` around lines 153 - 156, In
fingerprint_dataset initialize the missing variables before the main
fingerprinting loop: set spatial_dims to None or 0 as a placeholder, set
class_voxel_counts to a zeros container (e.g., list of zeros or
defaultdict(int)) sized for num_classes, set total_voxels to 0, and set
all_images_per_modality to an empty list or dict (e.g., list of lists or
defaultdict(list)) depending on how images are aggregated; update subsequent
code to use these containers (functions/variables referenced:
fingerprint_dataset, spatial_dims, class_voxel_counts, total_voxels,
all_images_per_modality) so no NameError occurs at runtime.
Comment on lines
+84
to
+104
| if self.nonzero_only: | ||
| if mask is not None: | ||
| fg = image[mask > 0] | ||
| elif self.use_mask_for_norm: | ||
| # If we were told to use a mask but none was provided, | ||
| # fall back to nonzero or raise? Official uses the seg provided to .run() | ||
| fg = image[image != 0] | ||
| else: | ||
| fg = image[image != 0] | ||
|
|
||
| if fg.size == 0: | ||
| return image # all-zero image: leave as-is | ||
|
|
||
| mean = float(fg.mean()) | ||
| std = max(float(fg.std()), 1e-8) | ||
| else: | ||
| mean = float(image.mean()) | ||
| std = max(float(image.std()), 1e-8) | ||
|
|
||
| image = (image - mean) / std | ||
| return image |
There was a problem hiding this comment.
Keep background voxels at zero when foreground-only stats are used.
mean and std are computed from foreground voxels, but the affine transform is then applied to the entire image. For masked or nonzero-only MRI inputs that turns zero background into a constant negative value, which defeats use_mask_for_norm.
🛠️ Proposed fix
- image = (image - mean) / std
- return image
+ if self.nonzero_only:
+ fg_mask = (mask > 0) if mask is not None else (image != 0)
+ out = np.zeros_like(image, dtype=np.float32)
+ out[fg_mask] = (image[fg_mask] - mean) / std
+ return out
+ return (image - mean) / std📝 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
| if self.nonzero_only: | |
| if mask is not None: | |
| fg = image[mask > 0] | |
| elif self.use_mask_for_norm: | |
| # If we were told to use a mask but none was provided, | |
| # fall back to nonzero or raise? Official uses the seg provided to .run() | |
| fg = image[image != 0] | |
| else: | |
| fg = image[image != 0] | |
| if fg.size == 0: | |
| return image # all-zero image: leave as-is | |
| mean = float(fg.mean()) | |
| std = max(float(fg.std()), 1e-8) | |
| else: | |
| mean = float(image.mean()) | |
| std = max(float(image.std()), 1e-8) | |
| image = (image - mean) / std | |
| return image | |
| if self.nonzero_only: | |
| if mask is not None: | |
| fg = image[mask > 0] | |
| elif self.use_mask_for_norm: | |
| # If we were told to use a mask but none was provided, | |
| # fall back to nonzero or raise? Official uses the seg provided to .run() | |
| fg = image[image != 0] | |
| else: | |
| fg = image[image != 0] | |
| if fg.size == 0: | |
| return image # all-zero image: leave as-is | |
| mean = float(fg.mean()) | |
| std = max(float(fg.std()), 1e-8) | |
| else: | |
| mean = float(image.mean()) | |
| std = max(float(image.std()), 1e-8) | |
| if self.nonzero_only: | |
| fg_mask = (mask > 0) if mask is not None else (image != 0) | |
| out = np.zeros_like(image, dtype=np.float32) | |
| out[fg_mask] = (image[fg_mask] - mean) / std | |
| return out | |
| return (image - mean) / std |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/dataloader/nnunet/normalization.py` around lines 84 - 104, The
current flow computes mean/std from foreground (fg) but applies the affine
transform to the whole image; change the logic in the normalization block
(referencing self.nonzero_only, mask, self.use_mask_for_norm, fg, mean, std,
image) so that when foreground-only stats are used you only normalize the
foreground voxels and leave background voxels as zero. Concretely, compute
mean/std from fg as before, then compute normalized_fg = (fg - mean) / std and
write those values back into the corresponding positions of image (using mask >
0 or image != 0 to find indices) and return image; when not using nonzero_only
keep the existing whole-image normalization path.
Comment on lines
+161
to
+172
| voxels_list = [] | ||
| for img in images: | ||
| flat = img.astype(np.float32).ravel() | ||
| if nonzero_only: | ||
| flat = flat[flat != 0] | ||
| voxels_list.append(flat) | ||
|
|
||
| if not voxels_list: | ||
| return {} | ||
|
|
||
| all_voxels = np.concatenate(voxels_list, axis=0) | ||
|
|
There was a problem hiding this comment.
compute_intensity_stats() can load the full dataset into RAM twice.
This function stores every flattened volume in voxels_list and then builds a second full copy via np.concatenate(). On typical 3D datasets that can OOM during fingerprinting before training even starts. A streaming mean/std plus sampled voxels for percentiles would scale much better.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/dataloader/nnunet/normalization.py` around lines 161 - 172,
compute_intensity_stats currently accumulates every flattened image in
voxels_list and then concatenates them, which duplicates memory usage; replace
this with a streaming/stateless approach: compute mean and variance online
(e.g., Welford's algorithm) over images in the loop to produce mean and std
without storing full arrays, and concurrently perform reservoir sampling (or
fixed-size random subsampling per image) to collect a bounded sample set used
only for percentile computations; update references to nonzero_only and images
handling inside compute_intensity_stats so you never call np.concatenate on all
voxels and only compute percentiles from the sampled array while returning
mean/std from the streaming aggregator.
Comment on lines
+1
to
+41
| import logging | ||
| import random | ||
| import subprocess | ||
| from pathlib import Path | ||
|
|
||
| from medicai.dataloader.nnunet.augmentations import ( | ||
| AugmentationConfig, | ||
| AugmentationPipeline, | ||
| ) | ||
| from medicai.dataloader.nnunet.dataset import nnUNetDataset | ||
| from medicai.dataloader.nnunet.dataset_fingerprint import fingerprint_dataset | ||
| from medicai.dataloader.nnunet.manifest import DatasetManifest | ||
| from medicai.dataloader.nnunet.preprocessing import preprocess_dataset | ||
| from medicai.models.nnunet.dynamic_unet import build_unet_from_plan | ||
| from medicai.trainer.nnunet.cross_validation import ( | ||
| generate_splits, | ||
| load_splits, | ||
| normalize_case_id, | ||
| save_splits, | ||
| ) | ||
| from medicai.trainer.nnunet.planning.planners import ( | ||
| nnUNetPlanner, | ||
| nnUNetPlannerResEncL, | ||
| nnUNetPlannerResEncM, | ||
| ) | ||
| from medicai.trainer.nnunet.training.trainer import nnUNetTrainer | ||
| from medicai.trainer.nnunet.utils.config import ( | ||
| DatasetFingerprint, | ||
| TrainingConfig, | ||
| nnUNetPlan, | ||
| ) | ||
| from medicai.trainer.nnunet.utils.io import ( | ||
| collapse_single_channel, | ||
| infer_spatial_dims, | ||
| load_medical_image, | ||
| load_npz, | ||
| normalize_layout, | ||
| save_medical_image, | ||
| ) | ||
| from medicai.utils.inference import sliding_window_inference | ||
|
|
There was a problem hiding this comment.
Critical: Missing numpy import will cause NameError at runtime.
The module uses np extensively in the predict() and _postprocess_prediction() methods (lines 319, 337-338, 358, 362, 367, 370, 375) but never imports numpy.
🐛 Proposed fix
import logging
import random
import subprocess
from pathlib import Path
+import numpy as np
+
from medicai.dataloader.nnunet.augmentations import (🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/trainer/nnunet/pipeline.py` around lines 1 - 41, The module is
missing the numpy import used by predict and _postprocess_prediction, causing
NameError; add an import for numpy (e.g., import numpy as np) at the top of the
file alongside the other imports so references to np in predict(...) and
_postprocess_prediction(...) resolve correctly.
Comment on lines
+80
to
+91
| def _metric_monitor_name(self): | ||
| """Return the correct monitor name depending on deep supervision.""" | ||
| if self.metrics and hasattr(self.metrics[0], "name"): | ||
| metric_name = self.metrics[0].name | ||
| else: | ||
| metric_name = "loss" | ||
| use_ds = ( | ||
| self.cfg.deep_supervision and self.net_cfg is not None and self.net_cfg.deep_supervision | ||
| ) | ||
| if use_ds: | ||
| return f"val_final_{metric_name}" | ||
| return f"val_{metric_name}" |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cd medicai && find . -name "trainer.py" -type fRepository: innat/medic-ai
Length of output: 94
🏁 Script executed:
cat -n medicai/trainer/nnunet/training/trainer.py | sed -n '230,260p'Repository: innat/medic-ai
Length of output: 1453
🏁 Script executed:
rg "iters_per_epoch|save_every_n_epochs" -A 2 -B 2Repository: innat/medic-ai
Length of output: 3726
🏁 Script executed:
rg "class TrainingConfig" -A 30Repository: innat/medic-ai
Length of output: 2047
🌐 Web query:
Keras ModelCheckpoint save_freq integer batches or epochs
💡 Result:
In tf.keras.callbacks.ModelCheckpoint, save_freq can be:
save_freq="epoch": saves at the end of every epoch. [1]save_freq=<int>: saves at the end of every N training batches (steps) — not epochs. [1]
Note: if you compiled the model with steps_per_execution=N, the “should I save now?” check only happens every N batches. [1]
Keras’ own guide materials describe the same behavior: "epoch" or an integer meaning “every N batches.” [2]
Sources: [1] TensorFlow ModelCheckpoint docs (save_freq definition) (tensorflow.org), [2] Keras guide noting save_freq is "epoch" or integer (every N batches) (keras.io)
Citations:
- 1: https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/ModelCheckpoint
- 2: https://keras.io/guides/orbax_checkpoint/?utm_source=openai
🌐 Web query:
Keras ModelCheckpoint mode="auto" loss metric
💡 Result:
In keras.callbacks.ModelCheckpoint, mode="auto" chooses whether “best” means minimize or maximize based on the name of the monitored quantity:
- If
monitorcontains"acc"(e.g.,val_accuracy,val_acc) or starts with"fmeasure", thenmode="auto"behaves likemode="max"(higher is better). - Otherwise (including
monitor="loss"ormonitor="val_loss"),mode="auto"behaves likemode="min"(lower is better). [1], [2]
So for a loss metric, mode="auto" will pick minimization (same as mode="min"). [1]
Sources: [1] [2]
🏁 Script executed:
rg "self.metrics" -B 5 -A 5 medicai/trainer/nnunet/training/trainer.pyRepository: innat/medic-ai
Length of output: 1070
🏁 Script executed:
rg "def _metric_monitor_name" -A 20 medicai/trainer/nnunet/training/trainer.pyRepository: innat/medic-ai
Length of output: 854
🏁 Script executed:
rg "val_every_n_epochs" medicai/trainer/nnunet/training/trainer.pyRepository: innat/medic-ai
Length of output: 40
The checkpoint callback setup has two issues with Keras API semantics.
-
save_frequses batches, not epochs: Whensave_freqis an integer inModelCheckpoint, Keras interprets it as "every N training batches", not epochs. Currentlysave_freq=cfg.save_every_n_epochs(default 50) saves every 50 batches, not every 50 epochs. This should multiply bycfg.iters_per_epochto achieve the intended behavior. -
mode="max"fails for loss metrics: When_metric_monitor_name()falls back toval_lossorval_final_loss(when no custom metrics are provided), the fixedmode="max"will keep worse (higher-loss) checkpoints instead of better ones. Usingmode="auto"allows Keras to infer the correct direction based on the metric name—mode="min"for loss,mode="max"for accuracy/dice metrics.
🛠️ Proposed fix
keras.callbacks.ModelCheckpoint(
filepath=str(self.output_dir / "best_model.weights.h5"),
monitor=monitor_name,
- mode="max",
+ mode="auto",
save_best_only=True,
save_weights_only=True,
verbose=1,
),
keras.callbacks.ModelCheckpoint(
filepath=str(self.output_dir / "checkpoint_latest.weights.h5"),
- save_freq=cfg.save_every_n_epochs,
+ save_freq=cfg.save_every_n_epochs * cfg.iters_per_epoch,
save_weights_only=True,
verbose=0,
),🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/trainer/nnunet/training/trainer.py` around lines 80 - 91, The
ModelCheckpoint setup is using save_freq as epochs and forcing mode="max", which
is incorrect; in the checkpoint creation (where ModelCheckpoint is instantiated)
multiply cfg.save_every_n_epochs by cfg.iters_per_epoch so save_freq =
cfg.save_every_n_epochs * cfg.iters_per_epoch (handle None/0 safely), and change
mode="max" to mode="auto" so Keras infers min/max from the metric name produced
by _metric_monitor_name (which returns "val_loss"/"val_final_loss" when no
metrics are present). Ensure references to cfg.save_every_n_epochs,
cfg.iters_per_epoch, ModelCheckpoint and _metric_monitor_name are updated
accordingly.
Comment on lines
+262
to
+268
| fit_history = self.model.fit( | ||
| self.train_dataset, | ||
| validation_data=self.val_dataset, | ||
| epochs=cfg.n_epochs, | ||
| steps_per_epoch=cfg.iters_per_epoch, | ||
| callbacks=all_callbacks, | ||
| verbose=1, |
There was a problem hiding this comment.
val_every_n_epochs is currently a dead config flag.
run() always validates every epoch because the fit call never uses cfg.val_every_n_epochs. On long 3D runs that adds avoidable validation cost and makes the config ineffective.
🛠️ Proposed fix
fit_history = self.model.fit(
self.train_dataset,
validation_data=self.val_dataset,
epochs=cfg.n_epochs,
steps_per_epoch=cfg.iters_per_epoch,
+ validation_freq=cfg.val_every_n_epochs,
callbacks=all_callbacks,
verbose=1,
)📝 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
| fit_history = self.model.fit( | |
| self.train_dataset, | |
| validation_data=self.val_dataset, | |
| epochs=cfg.n_epochs, | |
| steps_per_epoch=cfg.iters_per_epoch, | |
| callbacks=all_callbacks, | |
| verbose=1, | |
| fit_history = self.model.fit( | |
| self.train_dataset, | |
| validation_data=self.val_dataset, | |
| epochs=cfg.n_epochs, | |
| steps_per_epoch=cfg.iters_per_epoch, | |
| validation_freq=cfg.val_every_n_epochs, | |
| callbacks=all_callbacks, | |
| verbose=1, |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/trainer/nnunet/training/trainer.py` around lines 262 - 268, The
training run currently always validates every epoch because self.model.fit is
called with validation_data=self.val_dataset and doesn't respect
cfg.val_every_n_epochs; update run() to pass validation_freq (or implement an
equivalent callback) so validation only runs every cfg.val_every_n_epochs epochs
when cfg.val_every_n_epochs > 1, referencing self.model.fit,
cfg.val_every_n_epochs, self.train_dataset and self.val_dataset to conditionally
set validation_freq (or wrap/replace validation with a callback that runs
validation every N epochs) and leave behavior unchanged when
cfg.val_every_n_epochs == 1 or None.
Comment on lines
+120
to
+141
| spatial_dims=3, | ||
| patch_size=None, | ||
| batch_size=2, | ||
| n_pooling=5, | ||
| base_filters=32, | ||
| max_filters=320, | ||
| kernel_size=None, | ||
| pool_op_kernel_sizes=None, | ||
| conv_per_stage=2, | ||
| deep_supervision=True, | ||
| n_classes=2, | ||
| n_modalities=1, | ||
| output_activation="softmax", | ||
| ): | ||
| self.spatial_dims = spatial_dims | ||
| self.patch_size = patch_size if patch_size is not None else [128, 128, 128] | ||
| self.batch_size = batch_size | ||
| self.n_pooling = n_pooling | ||
| self.base_filters = base_filters | ||
| self.max_filters = max_filters | ||
| self.kernel_size = kernel_size if kernel_size is not None else [3, 3, 3] | ||
| self.pool_op_kernel_sizes = pool_op_kernel_sizes if pool_op_kernel_sizes is not None else [] |
There was a problem hiding this comment.
Make the default shapes depend on spatial_dims.
NetworkConfig(spatial_dims=2) currently defaults to patch_size=[128, 128, 128] and kernel_size=[3, 3, 3]. That leaks invalid 3D shapes into the 2D plan path whenever a caller relies on defaults.
🛠️ Proposed fix
self.spatial_dims = spatial_dims
- self.patch_size = patch_size if patch_size is not None else [128, 128, 128]
+ default_patch_size = [128] * spatial_dims
+ default_kernel_size = [3] * spatial_dims
+ self.patch_size = patch_size if patch_size is not None else default_patch_size
self.batch_size = batch_size
self.n_pooling = n_pooling
self.base_filters = base_filters
self.max_filters = max_filters
- self.kernel_size = kernel_size if kernel_size is not None else [3, 3, 3]
+ self.kernel_size = kernel_size if kernel_size is not None else default_kernel_size📝 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
| spatial_dims=3, | |
| patch_size=None, | |
| batch_size=2, | |
| n_pooling=5, | |
| base_filters=32, | |
| max_filters=320, | |
| kernel_size=None, | |
| pool_op_kernel_sizes=None, | |
| conv_per_stage=2, | |
| deep_supervision=True, | |
| n_classes=2, | |
| n_modalities=1, | |
| output_activation="softmax", | |
| ): | |
| self.spatial_dims = spatial_dims | |
| self.patch_size = patch_size if patch_size is not None else [128, 128, 128] | |
| self.batch_size = batch_size | |
| self.n_pooling = n_pooling | |
| self.base_filters = base_filters | |
| self.max_filters = max_filters | |
| self.kernel_size = kernel_size if kernel_size is not None else [3, 3, 3] | |
| self.pool_op_kernel_sizes = pool_op_kernel_sizes if pool_op_kernel_sizes is not None else [] | |
| spatial_dims=3, | |
| patch_size=None, | |
| batch_size=2, | |
| n_pooling=5, | |
| base_filters=32, | |
| max_filters=320, | |
| kernel_size=None, | |
| pool_op_kernel_sizes=None, | |
| conv_per_stage=2, | |
| deep_supervision=True, | |
| n_classes=2, | |
| n_modalities=1, | |
| output_activation="softmax", | |
| ): | |
| self.spatial_dims = spatial_dims | |
| default_patch_size = [128] * spatial_dims | |
| default_kernel_size = [3] * spatial_dims | |
| self.patch_size = patch_size if patch_size is not None else default_patch_size | |
| self.batch_size = batch_size | |
| self.n_pooling = n_pooling | |
| self.base_filters = base_filters | |
| self.max_filters = max_filters | |
| self.kernel_size = kernel_size if kernel_size is not None else default_kernel_size | |
| self.pool_op_kernel_sizes = pool_op_kernel_sizes if pool_op_kernel_sizes is not None else [] |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/trainer/nnunet/utils/config.py` around lines 120 - 141, The defaults
for 2D/3D are hardcoded to 3D sizes; update the constructor of the
NetworkConfig-like class to choose defaults based on spatial_dims: set
patch_size to [128]*spatial_dims when patch_size is None, kernel_size to
[3]*spatial_dims when kernel_size is None, and set pool_op_kernel_sizes to a
sensible default matching spatial_dims (e.g., empty list or list of kernel
tuples sized by spatial_dims) when pool_op_kernel_sizes is None; modify the
initialization logic around the attributes patch_size, kernel_size, and
pool_op_kernel_sizes in the __init__ (the block referencing self.patch_size,
self.kernel_size, self.pool_op_kernel_sizes) so callers with spatial_dims=2 get
2D shapes and spatial_dims=3 get 3D shapes.
Comment on lines
+48
to
+49
| elif image.ndim == 4: | ||
| data_out = image.transpose(3, 2, 1, 0).astype(dtype) |
There was a problem hiding this comment.
The 4D save transpose is not the inverse of load_nifti().
load_nifti() maps [X, Y, Z, C] to [Z, Y, X, C] with (2, 1, 0, 3). The inverse here should therefore also be (2, 1, 0, 3), but the current (3, 2, 1, 0) writes [C, X, Y, Z] and corrupts multi-channel NIfTI outputs.
🛠️ Proposed fix
- elif image.ndim == 4:
- data_out = image.transpose(3, 2, 1, 0).astype(dtype)
+ elif image.ndim == 4:
+ data_out = image.transpose(2, 1, 0, 3).astype(dtype)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@medicai/trainer/nnunet/utils/io.py` around lines 48 - 49, The 4D write path
in utils/io.py uses image.transpose(3,2,1,0) which is not the inverse of
load_nifti(); change the transpose axes in the image.ndim == 4 branch so that
save mirrors load_nifti()'s (2,1,0,3) mapping (i.e., reverse the same axes order
used in load_nifti()) to preserve multi-channel NIfTI layout for the function
that handles saving (the branch handling image.ndim == 4).
innat
marked this pull request as draft
April 27, 2026 15:51
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
fix #30
Summary by CodeRabbit
New Features
Documentation