Skip to content

README.md

Latest commit

 

History

History
181 lines (141 loc) · 6.74 KB

README.md

File metadata and controls

181 lines (141 loc) · 6.74 KB
Raw

Three-axis Zebrafinch nucleus segmentation with TriSAM

This tutorial shows how to segment nuclei in a 3D EM volume with a two-stage TriSAM workflow:

  1. Automatic proposals: microSAM's EM-organelles decoder generates approximate 3D nucleus instances.
  2. Three-axis refinement: vanilla SAM propagates a prompt from every proposal independently through XY, XZ, and YZ. Strict intersection of the three 3D masks is used as the final instance segmentation.

The second stage is genuine three-axis prompting; it does not use TriSAM's single best plane shortcut.

Requirements

Install a CUDA-compatible PyTorch build, then install the repository requirements:

python -m pip install -r requirements-zebrafinch.txt

Download an official SAM checkpoint, for example sam_vit_b_01ec64.pth. microSAM downloads the registered vit_b_em_organelles checkpoint automatically on first use.

The scripts require one CUDA GPU by default. Pass --device cpu only for small debugging runs; three-axis inference is otherwise impractically slow.

Input format

The example expects two HDF5 files with the same 3D shape and dataset key main:

image.h5   main: uint8 EM image, shape (Z, Y, X)
nuclei.h5  main: integer instance labels, shape (Z, Y, X); 0 is background

The label volume is used for evaluation and for the Zebrafinch proofreading exclusions. With --seed-source ais, labels do not provide the prompts for three-axis SAM.

If your annotation file uses a key other than main, pass --key YOUR_KEY to both stages.

Start with an ROI that contains several complete nuclei. The commands below use z=32:288, y=588:972, x=61:445 only as an example.

Non-interactive Slurm script

Run this from the TriSAM repository root. Change the paths and Slurm partition for your system.

#!/usr/bin/env bash
set -euo pipefail

IMAGE=/path/to/image.h5
GT=/path/to/nuclei.h5
SAM_CHECKPOINT=/path/to/sam_vit_b_01ec64.pth
RUN_DIR=/path/to/output/nuclei-example
PYTHON_BIN=/path/to/python

mkdir -p "$RUN_DIR/embeddings" logs
export TRISAM_NUCLEI_PYTHON="$PYTHON_BIN"

# Stage 1: automatic microSAM proposals. Invert the EM contrast when nuclei
# are bright under the original polarity.
AIS_JOB=$(sbatch --parsable --partition=short \
  scripts/run_zebrafinch_nuclei.sbatch \
  --image "$IMAGE" \
  --gt "$GT" \
  --key main \
  --roi-start 32 588 61 \
  --roi-shape 256 384 384 \
  --model-type vit_b_em_organelles \
  --invert \
  --foreground-threshold 0.01 \
  --center-distance-threshold 0.3 \
  --boundary-distance-threshold 0.7 \
  --embedding-path "$RUN_DIR/embeddings/microsam.zarr" \
  --metrics "$RUN_DIR/microsam-ais-metrics.json" \
  --output "$RUN_DIR/microsam-ais.h5")

# Stage 2: use every AIS instance as an automatic prompt, track it in all
# three orientations, and keep only voxels supported by XY, XZ, and YZ.
TRISAM_JOB=$(sbatch --parsable --partition=short \
  --dependency="afterok:${AIS_JOB}" \
  scripts/run_zebrafinch_nuclei_trisam.sbatch \
  --image "$IMAGE" \
  --gt "$GT" \
  --key main \
  --roi-start 32 588 61 \
  --roi-shape 256 384 384 \
  --seed-source ais \
  --seed-prediction "$RUN_DIR/microsam-ais.h5" \
  --model-type vit_b \
  --checkpoint "$SAM_CHECKPOINT" \
  --tau 0.80 \
  --prompt point+bbox \
  --primary-fusion intersection \
  --metrics "$RUN_DIR/trisam-three-axis-metrics.json" \
  --output "$RUN_DIR/trisam-three-axis.h5")

echo "microSAM job: $AIS_JOB"
echo "three-axis TriSAM job: $TRISAM_JOB"

Monitor the jobs with:

squeue -j "${AIS_JOB},${TRISAM_JOB}" \
  -o '%.18i %.28j %.12T %.10M %.10l %.6D %R'

Run without Slurm

On a GPU workstation, run the same entry points directly:

python model/zebrafinch_nuclei.py \
  --image /path/to/image.h5 \
  --gt /path/to/nuclei.h5 \
  --roi-start 32 588 61 \
  --roi-shape 256 384 384 \
  --invert \
  --foreground-threshold 0.01 \
  --embedding-path outputs/nuclei-example/embeddings/microsam.zarr \
  --metrics outputs/nuclei-example/microsam-ais-metrics.json \
  --output outputs/nuclei-example/microsam-ais.h5

python model/zebrafinch_nuclei_trisam.py \
  --image /path/to/image.h5 \
  --gt /path/to/nuclei.h5 \
  --roi-start 32 588 61 \
  --roi-shape 256 384 384 \
  --seed-source ais \
  --seed-prediction outputs/nuclei-example/microsam-ais.h5 \
  --checkpoint /path/to/sam_vit_b_01ec64.pth \
  --tau 0.80 \
  --primary-fusion intersection \
  --output outputs/nuclei-example/trisam-three-axis.h5

Output

The stage-2 HDF5 file contains:

image                    input ROI
gt_raw                   original labels in the ROI
gt_clean                 labels after Zebrafinch proofreading rules
prediction               selected uint32 3D instances; 0 is background
prediction_axis0         XY propagation
prediction_axis1         XZ propagation
prediction_axis2         YZ propagation
prediction_union         at least one view
prediction_majority      at least two views
prediction_intersection  all three views
prediction_robust        drift-filtered adaptive voting

Metrics and all run parameters are stored in the JSON metrics file and in the HDF5 metadata_json attribute.

Optional visualization

python scripts/visualize_nucleus_trisam_comparison.py \
  --baseline outputs/nuclei-example/microsam-ais.h5 \
  --trisam outputs/nuclei-example/trisam-three-axis.h5 \
  --output outputs/nuclei-example/comparison.png

Choosing the fusion rule

  • intersection is the recommended automatic setting. It rejects false tracks that appear in only one or two orientations.
  • majority has higher recall but can retain duplicate or non-nucleus prompts.
  • robust is useful for boundary-heavy ROIs where a cropped nucleus may be missing from one view.
  • union is usually too permissive for compact nuclei.

Limitations

  • The current evaluation path expects an instance-label volume, even when AIS supplies all prompts.
  • The proofreading exclusions in zebrafinch_nuclei.py are Zebrafinch-specific. Remove or replace them for another dataset.
  • Begin with ROIs. Full-volume inference still needs overlapping tiles, cross-tile instance deduplication, and rejection of vessel-lumen or other non-nucleus proposals.
  • Apparent precision is unreliable when the evaluation labels are incomplete.

Relevant code