feat: switch k-fold splitting to StratifiedGroupKFold for slide-disjoint folds#12
Conversation
…joint folds Replace StratifiedKFold with StratifiedGroupKFold so all tiles from a single slide land in the same fold. Rare classes (fewer than n_folds distinct slides) are now dropped entirely rather than collapsed to background. Removes the derived label column from output parquet since training re-labels via threshold. Adds per-fold tile percentage and fold-size CV metrics. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
📝 WalkthroughWalkthroughAdds a new ChangesK-fold splitting refactor with group-aware stratification and rare-class filtering
Sequence DiagramsequenceDiagram
participant Main as main()
participant RareFilter as drop_rare_class_slides()
participant Assign as assign_stratified_group_folds()
participant Log as log_fold_statistics()
Main->>RareFilter: labels, slide_ids, n_folds
RareFilter-->>Main: keep_mask
Main->>Main: filter labels, tissue_props, slide_ids, dataset via keep_mask
Main->>Assign: filtered labels, groups=slide_ids, n_folds, random_state
Assign-->>Main: folds
Main->>Log: labels, stratification_labels?, tissue_props, slide_ids, folds, n_folds
Log-->>Main: logged metrics (per-fold counts, tissue_prop stats, fold_size_cv)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
split/kfold_split.py (1)
43-50: ⚡ Quick winMake the rare-class filter contract match the actual mask semantics.
These lines describe dropping whole slides by dominant label, but Line 64 drops individual tile rows whose
labelis rare. Please update the docstring/function name, or change the mask to operate at slide granularity, so callers don't reason about the wrong filtering behavior.🤖 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 `@split/kfold_split.py` around lines 43 - 50, The docstring claims the function drops whole slides by dominant label but the implementation currently drops individual tile rows whose label is rare (uses the tile-level "label" mask); fix this mismatch by either (A) updating the docstring/function name to state it filters rare tile labels, or (B) changing the mask logic to compute each slide's dominant label (group by slide id, e.g., "slide" or "slide_id"), identify classes with fewer than n_folds slides, and produce a boolean keep-mask that sets all tiles from those rare slides to False so whole slides are dropped as described; adjust any references to StratifiedGroupKFold accordingly.
🤖 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 `@split/kfold_split.py`:
- Around line 43-50: The docstring claims the function drops whole slides by
dominant label but the implementation currently drops individual tile rows whose
label is rare (uses the tile-level "label" mask); fix this mismatch by either
(A) updating the docstring/function name to state it filters rare tile labels,
or (B) changing the mask logic to compute each slide's dominant label (group by
slide id, e.g., "slide" or "slide_id"), identify classes with fewer than n_folds
slides, and produce a boolean keep-mask that sets all tiles from those rare
slides to False so whole slides are dropped as described; adjust any references
to StratifiedGroupKFold accordingly.
![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 transitions the data splitting logic from StratifiedKFold to StratifiedGroupKFold, ensuring that all tiles from a specific slide are assigned to the same fold. It replaces the previous rare-class collapsing strategy with a filtering approach that drops tiles belonging to classes present in fewer than n_folds slides. Additionally, the PR enhances logging by including fold size coefficient of variation and tile percentages. Review feedback highlights a discrepancy in the drop_rare_class_slides docstring, which incorrectly refers to 'dominant labels' instead of the actual label-frequency filtering implementation.
3 participants
Problem
The existing
StratifiedKFoldsplit operates at tile level. Tiles from the sameslide can therefore land in both train and val within a single fold, which can
cause leakage when validation patches come from slides seen during training.
At the same time, the original tile-level split is still useful as a baseline and
should remain available for comparison.
Solution
Keep the original
StratifiedKFoldimplementation and add a configurableStratifiedGroupKFoldoption usingslide_idas the grouping key.The split strategy is selected through Hydra config:
split/kfold_split_5_folds: original tile-levelStratifiedKFoldsplit/stratified_group_kfold_split_5_folds: slide-groupedStratifiedGroupKFoldChanges
split/kfold_split.py: supports bothstratifiedandstratified_groupstrategies via
kfold_strategystratified: rare tile labels arecollapsed to
backgroundfor stratification onlystratified_group: labels appearing infewer than
n_foldsdistinct slides are dropped; count is logged asdropped_rare_class_tilesconfigs/experiment/split/stratified_group_kfold_split_5_folds.yamlkfold_strategyto split config metadata/hyperparametersscripts/submit_kfold_split.pyso the submitted approach is selectedby changing the experiment config
labelout of the output parquet; it is only an internal stratificationproxy, while downstream training derives labels from
roi_coverage_*fold_{i}_val_tile_pct,fold_size_cvResults (5-fold run, 1 102 086 tiles, 137 slides)
Fold size CV: 0.052. No rare-class drops in this dataset.
Summary by CodeRabbit
New Features
Improvements
Bug Fixes