Skip to content

Perf: vectorize _get_pairwise_iou with a joint bincount#9004

Open
aymuos15 wants to merge 2 commits into
Project-MONAI:devfrom
aymuos15:perf/panoptic-pairwise-iou-bincount
Open

Perf: vectorize _get_pairwise_iou with a joint bincount#9004
aymuos15 wants to merge 2 commits into
Project-MONAI:devfrom
aymuos15:perf/panoptic-pairwise-iou-bincount

Conversation

@aymuos15

@aymuos15 aymuos15 commented Jul 18, 2026

Copy link
Copy Markdown
Contributor

Description

Replaces the nested per-pair loop in _get_pairwise_iou (which held T + P full-volume masks live
and did O(pairs * volume) work) with a single joint bincount over gt * (num_pred + 1) + pred,
computing all pairwise intersections in one O(volume) pass that is much faster and, for realistic
instance counts, lower in peak memory.

Benchmark

Synthetic nuclei-style masks (random blobs), CPU, OMP_NUM_THREADS=8, median of repeated runs.

2D

Input (H×W, #true=#pred) Speedup Peak-mem saved
128×128, 40 ~11× 83%
256×256, 80 ~19.5× 95%
384×384, 120 ~19.9× 95%
512×512, 50 ~6× 96%
512×512, 150 ~18.3× 97%
512×512, 300 ~44.8× 98%
640×640, 180 ~21.8× 97%
768×768, 200 ~23.1× 96%
1024×1024, 150 ~16.1× 95%

3D

Input (D×H×W, #true=#pred) Speedup Peak-mem saved
64×64×64, 30 ~4.4× 92%
96×96×96, 50 ~4.6× 86%
128×128×128, 80 ~7.3× 93%
128×128×128, 10 ~1.2× 37%
160×160×160, 120 ~10× 95%

Numerics are bit-identical (torch.equal) at every size.

System: 12th Gen Intel Core i7-12800H (20 threads, OMP_NUM_THREADS=8); Linux 6.8.0-124-generic
x86_64, glibc 2.35; Python 3.10.12; PyTorch 2.12.1+cu130. Metric runs on CPU.

Types of changes

  • Non-breaking change (fix or new feature that would not break existing functionality).
  • New tests added to cover the changes.
  • The panoptic-quality metric and handler tests passed locally, along with black/isort/ruff/mypy.

Replace the nested per-pair loop in _get_pairwise_iou, which materialized one
full-volume mask per true and per predicted instance and computed full-volume
sum operations for every candidate pair, with a single joint bincount over
gt * (num_pred + 1) + pred. This counts all pairwise intersections and
per-instance areas in one O(volume) pass and derives IoU as
inter / (area_t + area_p - inter), the same integer-count formula the loop used.

Instance ids are contiguous after remap_instance_id, which both the old loop
(indexing mask lists by id) and this vectorized form rely on. Numerics are
bit-identical to the previous implementation and peak memory is lower, since
only a few bincount vectors and one (T+1)x(P+1) matrix are allocated instead
of T+P full-volume masks held simultaneously.

Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
@coderabbitai

coderabbitai Bot commented Jul 18, 2026

Copy link
Copy Markdown
Contributor
📝 Walkthrough

Walkthrough

_get_pairwise_iou now computes the IoU matrix with vectorized tensor operations, including early handling for empty instance sets and torch.bincount-based intersection and area calculations. Tests cover overlapping, disjoint, and empty ground-truth or prediction inputs.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The title clearly summarizes the main change: vectorizing _get_pairwise_iou with a joint bincount.
Description check ✅ Passed The description includes a clear summary, benchmark data, and type/test status, matching the template well enough.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🧹 Nitpick comments (1)
monai/metrics/panoptic_quality.py (1)

259-264: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Document the contiguous-ID precondition.

The bincount reshape relies on IDs being contiguous from zero; sparse IDs can produce an oversized bincount and fail the reshape. Add a Google-style docstring describing inputs, output, and this precondition.

As per path instructions, “Docstrings should be present for all definition” using Google-style sections.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@monai/metrics/panoptic_quality.py` around lines 259 - 264, Document the
function containing the gt_flat/pred_flat intersection computation with a
Google-style docstring covering its inputs, returned output, and the requirement
that remap_instance_id produces contiguous IDs starting at zero for both ground
truth and predictions. Ensure the docstring follows the project’s requirement
that every definition has documentation.

Source: Path instructions

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@monai/metrics/panoptic_quality.py`:
- Around line 259-264: Document the function containing the gt_flat/pred_flat
intersection computation with a Google-style docstring covering its inputs,
returned output, and the requirement that remap_instance_id produces contiguous
IDs starting at zero for both ground truth and predictions. Ensure the docstring
follows the project’s requirement that every definition has documentation.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 65a7dd41-409c-436e-8620-91cb89109f33

📥 Commits

Reviewing files that changed from the base of the PR and between 5da2472 and f443085.

📒 Files selected for processing (2)
  • monai/metrics/panoptic_quality.py
  • tests/metrics/test_compute_panoptic_quality.py

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant