lantern-ki3-abeta

LANTERN — Aβ plaque segmentation in 3D light-sheet microscopy (ki3 model 2)

5-fold ensemble of 3D U-Nets that segment amyloid-β plaques in cleared whole mouse brain imaged by light-sheet microscopy. Each fold is a frozen self-supervised ResEncL encoder with a fine-tuned decoder.

LANTERN = Light-sheet Automask Network for Transferable Embedding Representations.

What this is

Intended use

Detecting and delineating Aβ plaques in App-NLGF mouse brain, Abeta-stained channel, acquired at 4x and resampled to the near-iso-4 grid above. It is a research model for plaque burden quantification, not a diagnostic tool.

The 5 folds are meant to be run together: predict with each, threshold at 0.5, and take votes >= 3. The per-fold disagreement (votes/5) is useful on its own — model–model agreement correlates 0.85 with out-of-fold accuracy, so it flags regions needing human review without any ground truth.

Usage

The model is unusable without matching its preprocessing exactly — it is scale- and intensity-normalisation sensitive (see Limitations).

import numpy as np, torch

def normalize(vol):                      # per-tile, exactly as in training
    lo, hi = np.percentile(vol, [0.5, 99.5])
    v = np.clip(vol, lo, hi).astype(np.float32)
    s = float(v.std())
    return (v - float(v.mean())) / (s if s > 1e-8 else 1.0)

# `ResEncLUNet` comes from the LANTERN repo (lantern/models/unet_seg.py)
nets = []
for k in range(5):
    ck = torch.load(f"fold{k}/seg_model.pt", map_location="cpu", weights_only=False)
    n = ResEncLUNet(num_classes=2, num_input_channels=1)
    n.load_state_dict(ck["state_dict"])
    nets.append(n.eval().cuda())

x = torch.from_numpy(normalize(tile))[None, None].cuda()   # tile: float32 (128,128,128)
votes = torch.zeros((128, 128, 128), dtype=torch.uint8, device="cuda")
with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
    for n in nets:
        votes += (torch.softmax(n(x).float(), 1)[0, 1] >= 0.5).to(torch.uint8)
mask = (votes >= 3)

For volumes larger than 128³, tile at stride 64 (50% overlap) and combine overlapping tiles with a max over votes — a plaque clipped at one tile edge is recovered by the tile that contains it whole. Benchmarked throughput on one L40S: 27.8 ms/tile at bf16, batch 8 (batching past 4 buys almost nothing; bf16 over fp32 is worth ~1.4×).

Training data

75 patches of 128³ from 30 App-NLGF mice (mouse_app_lecanemab_ki3_aggregated), Abeta channel, hand-QC'd. Labels are re-curated: the previous model's held-out predictions were reviewed patch by patch and each patch assigned the segmentation judged best.

Ground truth carries a 9-voxel minimum connected-component floor, applied before training and before evaluation — objects smaller than that are neither learned nor scored.

Loss DiceCE with batch-Dice, 250 epochs, AdamW lr 1e-3, batch 2 × 4 accumulation, bf16. Augmentation: axis flips, one 90° rotation, mild gamma and Gaussian noise.

Evaluation

Object F1 at IoU ≥ 0.5 on 26-connected components, micro-averaged. Two references, and the difference matters:

The first row is inflated and should not be quoted as accuracy: 57 of 75 training labels are the previous model's own output, so the model is partly graded against an answer key its predecessor wrote. Against the fixed hand GT — the one reference that did not move — this model beats its predecessor by +0.019 object F1 (paired per-patch, better on 44/80). The real gain is precision, 0.380 → 0.413, at unchanged recall: fewer false positives, not more plaques found.

Voxel Dice is 0.806 but is a poor primary metric here — plaque foreground is ~1% of a patch, and Dice barely moves when ~100 false-positive components per patch are removed.

Limitations

• Scale sensitivity is severe. The model is trained at 4.4 × 3.25 × 3.25 µm and has no scale invariance. On a cohort acquired at z = 2.75 µm (1.6× finer) it predicted 0.21× the expected foreground volume. Resample to the training grid before inference; do not run it on a different pyramid level.
• Intensity normalisation must match. Percentile clip then z-score, computed per tile on raw intensities. The model was never trained on N4-corrected or otherwise rescaled data.
• Self-training bias. Most labels derive from an earlier model, so the ensemble reinforces that model's systematic biases. Judge it on human-annotated data only.
• Optimistic validation. The 5-fold split is patch-level, not subject-grouped, so patches from the same mouse appear in train and val. Expect worse performance on an unseen animal.
• Cohort narrowness. One study, App-NLGF genotype, one stain and acquisition protocol. Four PBS-treated subjects contributed no training patches, so the treated (Lecanemab) arm is over-represented: 41 patches vs 29.
• Small objects. Anything under 9 voxels is outside the task definition by construction.
• Not validated for absolute burden. Reported numbers are detection metrics on 128³ patches; whole-brain plaque-load quantification has not been validated against an independent measure.

Files

fold0..fold4/seg_model.pt   torch checkpoint: state_dict, config, per-fold val metrics
config.json                 architecture + preprocessing contract

Each checkpoint also carries the training config and that fold's validation metrics.

Citation

Code: https://github.com/Arshya-Guru/LANTERN

Released for non-commercial research use (CC BY-NC 4.0). The underlying imaging data is not included and is not released here.