Add optional RankSEG decoding to AsDiscrete#8908
Conversation
Signed-off-by: Zixun Wang <craddywang@gmail.com>
Signed-off-by: Zixun Wang <craddywang@gmail.com>
Signed-off-by: LI Junxing <83260380+Leev1s@users.noreply.github.com>
Signed-off-by: Zixun Wang <craddywang@gmail.com>
Signed-off-by: statmlben <bdai.hk@protonmail.com>
Signed-off-by: Zixun Wang <craddywang@gmail.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds optional RankSEG decoding to Estimated code review effort: 3 (Moderate) | ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
monai/transforms/post/dictionary.py (1)
198-200: ⚡ Quick winAdd
strict=Trueto zip() for robustness.Since the project requires Python ≥ 3.10, use
zip(..., strict=True)to catch configuration errors early whenargmaxandranksegsequences have mismatched lengths.🔒 Proposed fix
self.argmax = ensure_tuple_rep(argmax, len(self.keys)) self.rankseg = ensure_tuple_rep(rankseg, len(self.keys)) - if any(argmax_ and rankseg_ for argmax_, rankseg_ in zip(self.argmax, self.rankseg)): + if any(argmax_ and rankseg_ for argmax_, rankseg_ in zip(self.argmax, self.rankseg, strict=True)): raise ValueError("`rankseg=True` is incompatible with `argmax=True`.")🤖 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/transforms/post/dictionary.py` around lines 198 - 200, The zip call combining self.argmax and self.rankseg should use strict=True to fail fast on length mismatches; update the check in the Dictionary-based transform (around ensure_tuple_rep and the subsequent if that zips self.argmax and self.rankseg) to call zip(self.argmax, self.rankseg, strict=True) so a configuration error is raised if their lengths differ.Source: Linters/SAST tools
🤖 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/transforms/post/dictionary.py`:
- Around line 198-200: The zip call combining self.argmax and self.rankseg
should use strict=True to fail fast on length mismatches; update the check in
the Dictionary-based transform (around ensure_tuple_rep and the subsequent if
that zips self.argmax and self.rankseg) to call zip(self.argmax, self.rankseg,
strict=True) so a configuration error is raised if their lengths differ.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: fd566a60-5ae0-431b-89cc-bc846b4edb2f
📒 Files selected for processing (7)
docs/source/installation.mdmonai/transforms/post/array.pymonai/transforms/post/dictionary.pyrequirements-dev.txtsetup.cfgtests/transforms/test_as_discrete.pytests/transforms/test_as_discreted.py
ZixunWang
changed the title
Signed-off-by: Zixun Wang <craddywang@gmail.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/transforms/test_as_discrete.py (2)
89-94: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick winAdd docstring per coding guidelines.
Test methods should include Google-style docstrings describing what is tested. As per coding guidelines, "Docstrings should be present for all definition which describe each variable, return value, and raised exception."
📝 Proposed docstring
def test_rankseg_argmax_incompatible(self): + """Test that rankseg=True raises ValueError when combined with argmax=True.""" with self.assertRaises(ValueError):🤖 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 `@tests/transforms/test_as_discrete.py` around lines 89 - 94, Add a Google-style docstring to the test method test_rankseg_argmax_incompatible describing what the test verifies (that AsDiscrete raises ValueError when argmax and rankseg are both requested), list any parameters (none), the return (None), and the raised exception (ValueError) — reference the AsDiscrete constructor and call sites used in the test to explain the two asserted failure cases; keep the docstring concise and follow existing test docstring formatting in the file.Source: Coding guidelines
96-99: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick winAdd docstring per coding guidelines.
Test method should include Google-style docstring describing what is tested. As per coding guidelines, "Docstrings should be present for all definition which describe each variable, return value, and raised exception."
📝 Proposed docstring
def test_rankseg_missing_dependency(self): + """Test that rankseg=True raises OptionalImportError when rankseg package is unavailable.""" with mock.patch("monai.transforms.post.array.has_rankseg", False):🤖 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 `@tests/transforms/test_as_discrete.py` around lines 96 - 99, Add a Google-style docstring to the test_rankseg_missing_dependency function that succinctly describes the test purpose (verifying that AsDiscrete(..., rankseg=True) raises OptionalImportError when the rankseg dependency is missing), documents parameters (none), the return value (None), and the raised exception (OptionalImportError), and mentions the mocked condition (monai.transforms.post.array.has_rankseg set to False); place the docstring immediately under the def test_rankseg_missing_dependency(self): line to satisfy the project's docstring guideline and reference AsDiscrete and OptionalImportError in the description.Source: Coding guidelines
🧹 Nitpick comments (2)
tests/transforms/test_as_discrete.py (2)
99-99: 💤 Low valueUse
p()wrapper for consistency.Line 99 uses a raw list. Wrap with
p()fromTEST_NDARRAYSto test across multiple array types consistently.♻️ Proposed change
+ for p in TEST_NDARRAYS: with mock.patch("monai.transforms.post.array.has_rankseg", False): with self.assertRaises(OptionalImportError): - AsDiscrete(rankseg=True)([[[0.3, 0.6]], [[0.7, 0.4]]]) + AsDiscrete(rankseg=True)(p([[[0.3, 0.6]], [[0.7, 0.4]]]))🤖 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 `@tests/transforms/test_as_discrete.py` at line 99, The test passes a raw Python list to AsDiscrete(rankseg=True) instead of wrapping it with the p() array factory from TEST_NDARRAYS; update the call AsDiscrete(rankseg=True)([[[0.3, 0.6]], [[0.7, 0.4]]]) to use the p(...) wrapper from TEST_NDARRAYS so the input is created across all test ndarray types (e.g., AsDiscrete(rankseg=True)(p([[[0.3,0.6]], [[0.7,0.4]]]))); keep the same shape and values and ensure p is the TEST_NDARRAYS factory used elsewhere in the file.
94-94: 💤 Low valueUse
p()wrapper for consistency.Line 94 uses a raw list instead of wrapping with
p()fromTEST_NDARRAYS. For consistency with the test suite's pattern of testing across multiple array types, wrap the input.♻️ Proposed change
+ for p in TEST_NDARRAYS: with self.assertRaises(ValueError): - AsDiscrete(argmax=True)([[[0.3, 0.6]], [[0.7, 0.4]]], rankseg=True) + AsDiscrete(argmax=True)(p([[[0.3, 0.6]], [[0.7, 0.4]]]), rankseg=True)🤖 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 `@tests/transforms/test_as_discrete.py` at line 94, The test call to AsDiscrete(argmax=True) uses a raw Python list instead of using the TEST_NDARRAYS wrapper; update the input to wrap the list with the p() factory so the test runs across array types (use p([[ [0.3,0.6] ], [[0.7,0.4] ]]) or equivalent) before passing to AsDiscrete(argmax=True)(..., rankseg=True); locate the call to AsDiscrete and replace the raw list argument with p(...) to maintain consistency with the rest of the test suite.
🤖 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.
Outside diff comments:
In `@tests/transforms/test_as_discrete.py`:
- Around line 89-94: Add a Google-style docstring to the test method
test_rankseg_argmax_incompatible describing what the test verifies (that
AsDiscrete raises ValueError when argmax and rankseg are both requested), list
any parameters (none), the return (None), and the raised exception (ValueError)
— reference the AsDiscrete constructor and call sites used in the test to
explain the two asserted failure cases; keep the docstring concise and follow
existing test docstring formatting in the file.
- Around line 96-99: Add a Google-style docstring to the
test_rankseg_missing_dependency function that succinctly describes the test
purpose (verifying that AsDiscrete(..., rankseg=True) raises OptionalImportError
when the rankseg dependency is missing), documents parameters (none), the return
value (None), and the raised exception (OptionalImportError), and mentions the
mocked condition (monai.transforms.post.array.has_rankseg set to False); place
the docstring immediately under the def test_rankseg_missing_dependency(self):
line to satisfy the project's docstring guideline and reference AsDiscrete and
OptionalImportError in the description.
---
Nitpick comments:
In `@tests/transforms/test_as_discrete.py`:
- Line 99: The test passes a raw Python list to AsDiscrete(rankseg=True) instead
of wrapping it with the p() array factory from TEST_NDARRAYS; update the call
AsDiscrete(rankseg=True)([[[0.3, 0.6]], [[0.7, 0.4]]]) to use the p(...) wrapper
from TEST_NDARRAYS so the input is created across all test ndarray types (e.g.,
AsDiscrete(rankseg=True)(p([[[0.3,0.6]], [[0.7,0.4]]]))); keep the same shape
and values and ensure p is the TEST_NDARRAYS factory used elsewhere in the file.
- Line 94: The test call to AsDiscrete(argmax=True) uses a raw Python list
instead of using the TEST_NDARRAYS wrapper; update the input to wrap the list
with the p() factory so the test runs across array types (use p([[ [0.3,0.6] ],
[[0.7,0.4] ]]) or equivalent) before passing to AsDiscrete(argmax=True)(...,
rankseg=True); locate the call to AsDiscrete and replace the raw list argument
with p(...) to maintain consistency with the rest of the test suite.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 645ba23a-cff2-41c2-b3a3-3cd897653df8
📒 Files selected for processing (3)
monai/transforms/post/dictionary.pytests/transforms/test_as_discrete.pytests/transforms/test_as_discreted.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/transforms/test_as_discreted.py
- monai/transforms/post/dictionary.py
|
Really cool! To me, the only changes would be adding a 3D test and extending out a metric arg so that people can choose something other than dice? |
Signed-off-by: Zixun Wang <craddywang@gmail.com>
|
@aymuos15 Thanks for your review and suggestions! I have added a 3D test case: [[[[0.3, 0.6]]], [[[0.7, 0.4]]]] # shape: [C, H, W, D] = [2, 1, 1, 2]RankSEG internally handles inputs with shape The |
|
Thanks! Yeah, just wanted to make sure the internal handle is obvious here and the metric is an actual choice. LGTM! |
|
Hi @ericspod @Nic-Ma @KumoLiu, gentle ping on this PR. @aymuos15's feedback has been addressed in The checks are green and the PR appears to be mergeable cleanly. Would one of you be able to take a code-owner review, or let us know if there are any remaining API concerns before merging? Happy to make follow-up changes if needed. |
…onal-asdiscrete Signed-off-by: Zixun Wang <craddywang@gmail.com>
|
Hi @ericspod @Nic-Ma @KumoLiu, gentle follow-up on this PR. We have synced the branch with the latest MONAI dev branch and addressed the previous feedback. The intended usage follows the existing probs = Activations(softmax=True, dim=1)(logits)
post_argmax = AsDiscrete(argmax=True)
post_rankseg = AsDiscrete(rankseg=True, metric="dice")
pred_argmax = post_argmax(probs)
pred_rankseg = post_rankseg(probs)If the AsDiscrete(rankseg=True) / AsDiscreted(..., rankseg=True) API shape looks acceptable, would one of you be able to take a code-owner review? If you prefer a different API direction, happy to adjust. |
ericspod
left a comment
ericspod
left a comment
There was a problem hiding this comment.
Thanks @ZixunWang, I've looked over the PR with only some minor comments. I'm hesitant to approve this now as we want to limit the number of externald dependencies if we can, there are too many now as it is for a variety of reasons. I would like to have more discussion from other maintainers first. We also need to resolve merge conflicts.
| Defaults to ``None``. | ||
| rounding: if not None, round the data according to the specified option, | ||
| available options: ["torchrounding"]. | ||
| rankseg: whether to apply RankSEG decoding. Requires installing the optional ``rankseg`` package. |
There was a problem hiding this comment.
Probably this is going to be the first place users will encounter any reference to RankSEG, there should be some description of what it does and what it's for. It would help to reduce the length of the existing text as well.
There was a problem hiding this comment.
Thanks for the suggestion. I shortened the rankseg documentation and added a compact description clarifying that RankSEG is an inference-time decoder that converts class probability maps into segmentation label maps while maximizing the expected samplewise Dice or IoU.
| if argmax and rankseg: | ||
| raise ValueError("`rankseg=True` is incompatible with `argmax=True`.") | ||
| self.argmax = argmax | ||
| self.rankseg = rankseg |
There was a problem hiding this comment.
If rankseg is True but the package isn't present, it would help to raise the OptionalImportError here as well.
There was a problem hiding this comment.
Thanks. I added the dependency check to AsDiscrete.__init__(), so AsDiscrete(rankseg=True) now raises OptionalImportError immediately when the rankseg package is unavailable. I retained the check in __call__() for cases where RankSEG is enabled per call, and added test coverage for both paths.
| rounding: if not None, round the data according to the specified option, | ||
| available options: ["torchrounding"]. it also can be a sequence of str or None, | ||
| each element corresponds to a key in ``keys``. | ||
| rankseg: whether to apply RankSEG decoding. Requires installing the optional ``rankseg`` package. |
There was a problem hiding this comment.
Similarly this is likely where new users will encounter this so a comment to see the description in AsDiscrete would help.
There was a problem hiding this comment.
Thanks. I shortened the AsDiscreted parameter documentation and added a cross-reference to AsDiscrete for the full RankSEG description, while retaining the dictionary-specific sequence behavior.
Signed-off-by: statmlben <bdai.hk@protonmail.com>
|
Thank you @ericspod for the careful review! The inline comments and merge conflict have now been addressed, and the checks are green. I understand the concern about adding another external dependency. To provide more concrete evidence for the value of this optional integration, we performed a reproducible comparison between the existing argmax post-processing and RankSEG using the same pretrained model probability outputs. Experimental setup
Quantitative results
Compared with argmax:
Qualitative examplesThe figure below shows the input, ground truth, argmax prediction, RankSEG (Dice) prediction, and RankSEG error overlay. Pancreas is cyan and tumor is orange; in the error overlay, white is true positive, red is false positive, cyan is false negative, and yellow denotes a foreground class mismatch. The cases were selected using the 25th, 50th, and 75th percentiles of the paired Dice improvement rather than selecting only the largest improvements. For each 3D volume, the displayed slice maximizes ground-truth foreground area, and the crop is based only on the ground-truth foreground bounds; neither rule depends on either decoder's result. 
The evaluation code, fixed configuration, per-case results, summary, and raw timing samples are available at: The intended MONAI usage remains a small change to the standard decollated post-processing workflow: post_pred = Compose(
[
Activations(softmax=True),
AsDiscrete(rankseg=True, metric="dice"),
]
)
predictions = [post_pred(x) for x in decollate_batch(logits)]RankSEG remains fully optional: For additional context, I maintain RankSEG and supervised (my students) on this work. RankSEG is a long-term project that we intend to actively maintain, and I would be happy to help address future compatibility or MONAI integration issues. I appreciate the maintainers taking the time to consider the dependency trade-off, and I am happy to adjust the packaging or integration boundary if a different approach would be preferable. |
5 participants
Related to #8857.
Description
This PR adds optional RankSEG decoding support to MONAI post-processing.
RankSEG is integrated as an extension of the existing discrete post-processing API:
rankseg=Falseremains the default, so existing behavior is unchanged.Types of changes
./runtests.sh -f -u --net --coverage../runtests.sh --quick --unittests --disttests.make htmlcommand in thedocs/folder.Changes
ranksegsupport toAsDiscreteandAsDiscreted.Many thanks to @statmlben and @Leev1s for their collaboration and helpful discussions on the RankSEG integration.
Feedback Requested
This PR is intended to start from a conservative integration point. We would especially appreciate feedback on:
AsDiscrete(rankseg=True)/AsDiscreted(..., rankseg=True)is the preferred API shape;argmax=Trueandrankseg=Trueshould raise aValueErroras implemented, or only emit a warning.