apooladi/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

Architecture	ResEncLUNet — ResEncL encoder (6 stages, features [32,64,128,256,320,320]) + UNetDecoder
Parameters	102.3 M total, 12.1 M trainable (encoder frozen)
Encoder init	nnBYOL3D self-supervised pretraining on ~50k near-iso-4 light-sheet patches
Input	1 channel, 128³ voxels, near-iso-4 grid (4.4 × 3.25 × 3.25 µm)
Output	2 classes (background, plaque)
Folds	5, from patch-level 5-fold CV — use all 5 and vote

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.

label source	patches
model prediction (previous LANTERN model, held-out)	57
original hand GT	10
hand-revised model prediction	5
prediction from a lower size-floor model	1
hand GT with a size floor applied	2

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:

reference	object F1	precision	recall	FP/patch
the re-curated labels (out-of-fold)	0.6345	0.574	0.710	106
fixed original hand GT (80 patches)	0.4482	0.413	0.490	158
previous model, same fixed GT	0.4295	0.380	0.493	181

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.