[FIX] patch sampling in get_bbox() method to ensure patches are always fully contained inside the volume when oversampling foreground

[sifaoso]

Hey all,

Thanks a lot for all the work that has gone into nnU-Net over the years!

I believe I found a bug in patch sampling when visualizing 3D CT batches (AMOS dataset), where some sampled patches contain slices outside the scan field-of-view (resulting in unexpected padding).

The issue occurs in the get_bbox() method when force_fg=True (foreground oversampling). In this case, a foreground class is first selected, then a random voxel from that class is chosen as the patch center. However, when constructing the patch, the resulting bounding box is only constrained with respect to the lower bounds (lbs) and is not checked against the upper bounds (ubs), which define the valid sampling range to keep the patch fully inside the volume (here:

nnUNet/nnunetv2/training/dataloading/data_loader.py

Line 145 in e73bbc7

bbox_lbs = [max(lbs[i], selected_voxel[i + 1] - self.patch_size[i] // 2) for i in range(dim)]

). As a consequence, some foreground patches may partially fall outside the valid image region and become padded.

Note that this issue does not occur when force_fg=False, since bbox_lbs is sampled directly within the valid interval [lbs, ubs].

The proposed fix explicitly constrains the sampled bounding box to remain within both lbs and ubs, which fixes the issue in my tests.

Below is a small GIF illustrating the problem (see the 3rd sample: some slices lie outside the scan field-of-view and are therefore padded unexpectedly).

[ghost]

🔍 nnU-Net Code Review

Summary

This PR fixes a boundary-check bug in get_bbox() where, during foreground oversampling (force_fg=True), the computed bbox_lbs was clamped against the lower bound (lbs) but not the upper bound (ubs). When a selected foreground voxel lies near the end of an axis, the patch could extend beyond the valid image region, causing unexpected zero-padding during training.

Contribution Guidelines

• Bug fix format: The PR describes the root cause clearly and includes a visual (GIF) showing the issue. However, CONTRIBUTING.md asks for a minimal reproducible example — a short Python/CLI script that triggers the bug would be more rigorous than a screenshot/GIF alone. This is a minor gap rather than a blocker.
• General applicability: The fix is not dataset-specific; the bug exists for any dataset where a foreground voxel coordinate is close to the image edge in any axis. Observed on AMOS but affects all datasets.
• Scope: Minimal, single-line change — no unnecessary complexity introduced.

Key Findings

• Correctness — the fix is correct. lbs[i] and ubs[i] define the valid sampling interval for bbox_lbs in each dimension (nnunetv2/training/dataloading/data_loader.py:93-94). The non-force_fg branch already samples uniformly within [lbs[i], ubs[i]] (line 99). The force_fg branch missed the upper-bound clamp, which the PR correctly adds:

  # Before
  bbox_lbs = [max(lbs[i], selected_voxel[i + 1] - self.patch_size[i] // 2) for i in range(dim)]
  # After
  bbox_lbs = [max(lbs[i], min(ubs[i], selected_voxel[i + 1] - self.patch_size[i] // 2)) for i in range(dim)]

• Foreground guarantee is preserved. When the clamping triggers (selected_voxel[i+1] - patch_size[i]//2 > ubs[i]), bbox_lbs[i] is set to ubs[i] and bbox_ubs[i] = ubs[i] + patch_size[i]. The selected foreground voxel index satisfies selected_voxel[i+1] <= data_shape[i] - 1 < bbox_ubs[i] and selected_voxel[i+1] > ubs[i] = bbox_lbs[i], so the foreground voxel is still inside the resulting patch.
• The code comment at line 142–143 says "Make sure it is within the bounds of lb and ub", but the original code only enforced lb. The comment was already misleading before this fix; it is now accurate. No update to the comment is strictly required, but it could be clarified for future readers.
• No impact on other paths. The force_fg=False + has_ignore branch and the random-fallback branch are untouched.

Minor Suggestions

• Consider adding a short code-level comment (or updating line 143) to make clear that min(ubs[i], ...) is the fix for the upper-bound case — future readers will wonder why both bounds are needed.
• A minimal Python test or assertion demonstrating that bbox_lbs[i] <= ubs[i] always holds after the fix would strengthen the PR per CONTRIBUTING.md, even if it's just added to the PR description.

[FabianIsensee]

Thanks for the report. Looks legit! I am in vacation right now and will look into it once I am back