Generalize TestTimeAugmentation to non-spatial predictions

[ytl0623]

Fixes #8276

Description

• Added a new argument apply_inverse_to_pred. Defaults to True to preserve backward compatibility. When set to False, it skips the inverse transformation step and aggregates the model predictions directly.
• Added a new unit test to simulate a classification task with spatial augmentation, verifying that the aggregation works correctly without spatial inversion.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[coderabbitai]

📝 Walkthrough

Walkthrough

A new boolean parameter apply_inverse_to_pred (default=True) is added to the TestTimeAugmentation class, enabling conditional application of inverse transformations to predictions. When set to False, raw predictions bypass inversion, supporting non-spatial prediction types like classification scores. The __call__ method now includes conditional logic to handle both pathways. A corresponding test method validates the feature for non-spatial outputs. A UNet architecture parameter is also updated from strides=(2, 2) to strides=(2,) in the test file.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title accurately summarizes the main objective: generalizing TestTimeAugmentation to support non-spatial predictions via the new apply_inverse_to_pred parameter.
Description check	✅ Passed	Description covers the key changes (new argument, backward compatibility, test addition) and references the linked issue (#8276). Template mostly followed with change types marked.
Linked Issues check	✅ Passed	PR directly addresses #8276 by introducing apply_inverse_to_pred parameter to skip inverse transforms for non-spatial predictions and adding test coverage for classification scenarios.
Out of Scope Changes check	✅ Passed	All changes align with the scope of #8276. The UNet stride modification in tests appears to be a minor adjustment for test compatibility and doesn't introduce unrelated functionality.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

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

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}