Add LANTERN 5-fold deep-learning plaque segmentation - #192
Open
Arshya-Guru wants to merge 4 commits into
Open
Conversation
Adds an opt-in alternative to the GMM/N4 plaque segmentation: the LANTERN
ResEncL-UNet ensemble (apooladi/lantern-ki3-abeta), selected with
--seg_method lantern.
The model reads raw intensities -- it was never trained on N4-corrected data
-- so this path bypasses the bias-field chain entirely and reads the SPIM
input directly, as run_vesselfm does. It runs at plaque_level (1), the grid it
was built for; the segmentation_level mask is that prediction upsampled by
exact voxel replication rather than native inference at level 0.
Two outputs: a level-1 float32 probability map holding the fraction of folds
voting foreground ({0, .2, .4, .6, .8, 1} for five folds, the same form as
LANTERN's own whole-brain probmask), and the level-0 binary mask valued 0/100
to match the segmentation.smk convention, so fieldfrac / regionprops / counts /
segstats / heatmaps / QC all consume it with no changes to those rules.
Writing the probability map unthresholded means the vote threshold can be
changed without re-running the GPU; binarize compares against the midpoint
between attainable values so float representation cannot flip a boundary voxel.
ZarrNii.segment() could not be used: it dispatches through da.map_blocks with
disjoint blocks, no block position and a hardcoded uint8 output, so it can
express neither the stride-64 tile overlap the model requires nor brain-mask
block skipping. da.map_overlap is driven directly instead, with a halo of
tile - stride, brain restriction via the desc-brain mask (the n4_biasfield
precedent), and chunks laid out so no chunk is shorter than the halo.
Like vessels, the method is resolved to a single stain rather than fanned out:
it is trained on amyloid-beta only, so it stays out of the generic seg_method
expansions and the cross-stain aggregation the way vessel_seg_method does.
Remaining stains fall back to the --seg_method default, read off the CLI
definition so it cannot drift. Verified: against a run that never mentions
lantern, the job plan differs only by the added lantern jobs.
The architecture is vendored (lantern_model.py) rather than depended on, and
loads the published checkpoints with zero missing/unexpected keys.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018WbbraRDoezJrijVaAsL97
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 8 changed files in this pull request and generated 3 comments.
Suppressed comments (4)
pyproject.toml:57
- This global PyPI dependency is also installed in the
dev-onlyenvironment, which has no runtime PyTorch dependency. The updated lock consequently resolves PyPI Torch 2.13 plus the CUDA 13 toolkit stack there (pixi.lock:1759-1761), making the lint-only environment download a large, unrelated GPU runtime. Scope this package to the runtime/GPU feature instead of the workspace-global PyPI dependencies.
dynamic-network-architectures = ">=0.4"
spimquant/workflow/scripts/lantern_plaques.py:241
- The new chunk planning, halo handling, and overlap aggregation have no regression tests, although similar standalone Dask processing is directly unit-tested in
tests/test_skeleton_graph_from_sdt.py:63-90. Add CPU tests covering short remainder chunks, block-boundary continuity, background skipping, and folds detecting the same voxel from different overlapping tiles.
def vote_map(data, brain, predict, tile, stride):
"""Lazily map `predict` over `data`, skipping blocks that carry no brain.
The halo is `tile - stride`, so every voxel in a block's core is covered by
at least one tile lying wholly inside the block that produced it -- which is
spimquant/workflow/scripts/binarize_lantern_plaques.py:47
- The vote map written by inference contains normalized fractions in
[0, 1], not raw counts in0..n_folds, so this validation error gives incorrect diagnostic information. Describe the valid threshold range instead.
f"{n_folds}-fold ensemble; the vote map holds counts in 0..{n_folds}"
spimquant/workflow/scripts/lantern_plaques.py:95
- CUDA availability does not imply bfloat16 support; pre-Ampere CUDA devices can reach this branch and fail when autocast selects
bfloat16. Gate mixed precision withtorch.cuda.is_bf16_supported()(or choose a supported dtype) so the documented generic GPU path remains usable.
torch.autocast(
device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"
),
Comment on lines
+207
to
+209
| seg_methods = [ | ||
| m for m in config["seg_method"] if m != config["plaque_seg_method"] | ||
| ] or list(config["parse_args"]["--seg_method"]["default"]) |
Comment on lines
+102
to
+106
| for j, (z, y, x) in enumerate(batch): | ||
| region = votes[z : z + tile, y : y + tile, x : x + tile] | ||
| # MAX, not sum: a plaque clipped at one tile edge is recovered by the | ||
| # tile that contains it whole. | ||
| np.maximum(region, batch_votes[j], out=region) |
Comment on lines
+605
to
+609
| inputs["spim"].expand( | ||
| bids( | ||
| root=root, | ||
| datatype="seg", | ||
| stain="{stain}", |
Member
|
Excellent, nice to see this integrated! I'm thinking about whether we want to future-proof this up-front for later models we train on things other than plaques. I.e. instead of calling this lantern_plaques, just call it lantern, and then have the configuration parameters all indexed under a top-level stain channel. Then when we have another model (e.g. for Iba1) could also include all those parameters there.. This PR also highlights some inflexibility we have in spimquant's seg_method selection and being able to set things for specific channels, I'll think more about that one.. |
Author
|
Yes I was also wondering about that for once we have IBA1 models trained. Also, I checked copilot's comments here and and it did correctly flag some issues like LANTERN not produces stats and Abeta being segmented twice (the fallback GMM happens for the Abeta channel too). I'll try to resolve these and reship |
Addresses the Copilot review on #192. The plaque method is stain-specific, so excluding it from `seg_methods` and letting every target rule expand `desc x stain` had two consequences: - With `--seg_method lantern` as the only method, the fallback to the CLI default expanded across all of stains_for_seg, so the amyloid channel was segmented twice -- once by LANTERN and once by the default GMM. A dry run confirms this scheduled a redundant gmmthresh, n4_biasfield and n4_pre_quant for that stain. - LANTERN itself appeared in none of the desc expansions, so it produced no regionprops, counts, segstats, QC or sidecars, and could not be compared against the method it is meant to replace. Replace the cross product with `stains_for_desc(desc)` and the `seg_expand` / `seg_expand_plain` helpers, which expand each method against its own stains. Cross-stain targets now use `coloc_seg_methods`, which drops single-stain methods, and `do_coloc` is derived from it. The cross-stain aggregations in segstats.smk and regionprops.smk become input functions of `wildcards.desc` for the same reason; this also removes the unguarded `stains_for_seg[0]`. LANTERN now produces the same downstream set as gmm+n3k1 apart from the coloc-* products, which are undefined for a single-stain method. Also from the same review: - Scope dynamic-network-architectures to the `runtime` feature. At workspace level it pulled a PyPI torch and CUDA stack into the lint-only `dev-only` environment, which has no other torch dependency. - Correct the plaque_vote_threshold error message, which described the vote map as raw counts after it changed to normalized fractions. Not changed: folds are still summed within a tile before tiles are combined with max. Those steps do not commute, but this is the order LANTERN's own infer_wholebrain.py uses and plaque_vote_threshold was calibrated under it. Documented in place so it is not "fixed" later by mistake. Verified by dry run against a synthetic BIDS tree: per-rule job counts for a run that never mentions lantern are byte-identical to the previous commit. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018WbbraRDoezJrijVaAsL97
Requesting only the plaque method on a single-stain dataset (the actual production invocation: --seg_method lantern --stains_for_seg Abeta) removes that stain from stains_for_seg_methods and leaves the fallback generic method covering nothing. seg_expand already skipped it, but all_seg_methods still listed it, so the stain-less cross-stain targets kept requesting its products. Those rules then received an empty input list and failed at runtime -- merge_indiv_and_coloc_segstats_tsv died on `merged.to_csv` with merged=None. Filter all_seg_methods to methods that actually segment something. The previous commit's verification used two stains, which left one for the generic method and hid the degenerate case. Verified by dry run on a synthetic BIDS tree: - single stain + lantern: zero desc-gmm+n3k1 targets scheduled (was 12 merge_indiv_and_coloc_segstats_tsv + 3 aggregate_regionprops_across_stains, all of which failed on the cluster) - two stains + lantern: unchanged, gmm covers Iba1 and lantern covers Abeta - no lantern: per-rule job counts still byte-identical to 8bdb3ec Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018WbbraRDoezJrijVaAsL97
With `gpu: 2` and no explicit task setting, the SLURM executor appends
`--ntasks-per-gpu=1` (submit_string.py:71-83). Combined with `--gpus=2` that
requests two tasks, so srun runs the entire job twice, each task pinned to one
device. Observed on the cluster: the log contains "loaded 5 folds" twice,
"localrule run_lantern_plaques" twice, and
WARNING: 2 GPUs requested but only 1 visible; using 1
twice. Both copies then infer over the whole volume and race to write the same
probseg store.
The executor already provides an escape hatch for this -- its own comment cites
python process management interfering with SLURM for pytorch -- and skips the
flag when the value is < 1. Set tasks_per_gpu=0, which yields
--gpus=2 --cpus-per-gpu=8
i.e. one task that sees both devices, which is what DevicePool was written for.
run_vesselfm is unaffected: gpu=1 with --ntasks-per-gpu=1 is a single task.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018WbbraRDoezJrijVaAsL97
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds an opt-in alternative to the GMM/N4 plaque segmentation
Leaves in the old GMM based method as the default, but using the
--seg_method lanternwill use the ki3 trained amyloid beta plaque detection model (version 2). For now,--seg_method lanternwill only work on the amyloid beta channel and other non-vessel channels should revert to the default seg_method.