Thayer Alshaabi1,2*, Daniel E. Milkie1, Gaoxiang Liu2, Cyna Shirazinejad2, Jason L. Hong2, Kemal Achour2, Frederik Görlitz2, Ana Milunovic-Jevtic2, Cat Simmons2, Ibrahim S. Abuzahriyeh2, Erin Hong2, Samara Erin Williams2, Nathanael Harrison2, Evan Huang2, Eun Seok Bae2, Alison N. Killilea2, David G. Drubin2, Ian A. Swinburne2, Srigokul Upadhyayula2,3,4*, Eric Betzig1,2,5*
High-resolution tissue imaging is often compromised by sample-induced optical aberrations that degrade resolution and contrast. While wavefront sensor-based adaptive optics (AO) can measure these aberrations, such hardware solutions are typically complex, expensive to implement, and slow when serially mapping spatially varying aberrations across large fields of view. Here, we introduce AOViFT (Adaptive Optical Vision Fourier Transformer)---a machine learning-based aberration sensing framework built around a 3D multistage Vision Transformer that operates on Fourier domain embeddings. AOViFT infers aberrations and restores diffraction-limited performance in puncta-labeled specimens with substantially reduced computational cost, training time, and memory footprint compared to conventional architectures or real-space networks. We validated AOViFT on live gene-edited zebrafish embryos, demonstrating its ability to correct spatially varying aberrations using either a deformable mirror or post-acquisition deconvolution. By eliminating the need for the guide star and wavefront sensing hardware and simplifying the experimental workflow, AOViFT lowers technical barriers for high-resolution volumetric microscopy across diverse biological samples.
Docker & Apptainer images
Our prebuilt image with Python, TensorFlow, and all packages installed for you.
docker pull ghcr.io/cell-observatory/aovift:main_tf_cuda_12_3
Running an image on a cluster typically requires an Apptainer version of the image, which can be generated by:
apptainer pull --force main_tf_cuda_12_3.sif docker://ghcr.io/cell-observatory/aovift:main_tf_cuda_12_3
git clone --recurse-submodules https://github.com/cell-observatory/aovift.git
To later update to the latest, greatest
git pull --recurse-submodules
Run tests/test_tensorflow.py
to make sure the installation works
pytest -s -v --disable-pytest-warnings --color=yes tests/test_tensorflow.py
Note: If you want to run a local version of the image, see the Dockerfile
Pretrained models
All pre-trained models can be downloaded from our pretrained models repository.
If you wish to download all of our models at once, you can use this link and extract the desired *.h5 file from the zip file.
Note: To run AOViFT pytests you need to download the following files:
- Examples directory that has a few example tif files and extract it to
aovift/examples
. - aovift-15-YuMB-lambda510.h5 and save it to
aovift/pretrained_models
directory.
The [src/python ao.py
](src/python ao.py) script provides a CLI
for running our models on a given 3D stack (.tif
file).
To run docker image, replace working-dir
with your local path for the repository.
docker run --network host -u 1000 --privileged -v working-dir/aovift:/app/aovift --env PYTHONUNBUFFERED=1 --pull missing -t -i --rm -w /app/aovift --ipc host --gpus all ghcr.io/cell-observatory/aovift:main_tf_cuda_12_3 bash
Running tests/test_datasets.py
will create a dataset of synthetic data for testing.
pytest -s -v --disable-pytest-warnings --color=yes tests/test_datasets.py
Note: If you just want to use our beads simulator without our models for your own projects, you can use our beads simulator repository.
Running tests/test_embeddings.py
will create a dataset of synthetic data for testing.
pytest -s -v --disable-pytest-warnings --color=yes tests/test_embeddings.py
tests/test_embeddings.py::test_fourier_embeddings Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.125_y_0.125_z_0.2_twd_simulator_False
Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.097_y_0.097_z_0.2_twd_simulator_False
PASSED
tests/test_embeddings.py::test_interpolate_embeddings Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.125_y_0.125_z_0.2_twd_simulator_False
Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.097_y_0.097_z_0.2_twd_simulator_False
PASSED
tests/test_embeddings.py::test_rolling_fourier_embeddings Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.125_y_0.125_z_0.2_twd_simulator_False
Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.097_y_0.097_z_0.2_twd_simulator_False
Preprocessing, 1 rois per tile, roi size, (64, 64, 64), stride length [64 64 64], throwing away [0 0 0] voxels: 100%|█████████████████████████████████████████████| 1/1 [00:00<00:00, 7.84it/s] 0.1s elapsed
Compute FFTs: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 50.52it/s] 0.0s elapsed
Remove interference patterns: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:01<00:00, 1.04s/it] 1.0s elapsed
PASSED
tests/test_embeddings.py::test_embeddings_with_digital_rotations Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.125_y_0.125_z_0.2_twd_simulator_False
Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.097_y_0.097_z_0.2_twd_simulator_False
Preprocessing, 1 rois per tile, roi size, (64, 64, 64), stride length [64 64 64], throwing away [0 0 0] voxels: 100%|█████████████████████████████████████████████| 1/1 [00:00<00:00, 7.85it/s] 0.1s elapsed
Compute FFTs: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 149.57it/s] 0.0s elapsed
Remove interference patterns: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:01<00:00, 1.05s/it] 1.0s elapsed
PASSED
To predict the wavefront for a small FOV, you can use the predict_sample
function:
/app/aovift$ pytest -s -v --disable-pytest-warnings --color=yes tests/test_ao.py -k test_predict_sample
tests/test_ao.py::test_predict_sample Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.125_y_0.125_z_0.2_twd_simulator_False
Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.097_y_0.097_z_0.2_twd_simulator_False
6/6 [==============================] - 8s 891ms/step
Evaluate predictions: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:03<00:00, 3.26s/it] 3.3s elapsed
PASSED
usage: ao.py predict_sample [-h] [--current_dm CURRENT_DM] [--prev PREV] [--lateral_voxel_size LATERAL_VOXEL_SIZE] [--axial_voxel_size AXIAL_VOXEL_SIZE] [--wavelength WAVELENGTH]
[--dm_damping_scalar DM_DAMPING_SCALAR] [--freq_strength_threshold FREQ_STRENGTH_THRESHOLD] [--prediction_threshold PREDICTION_THRESHOLD]
[--confidence_threshold CONFIDENCE_THRESHOLD] [--sign_threshold SIGN_THRESHOLD] [--plot] [--plot_rotations] [--num_predictions NUM_PREDICTIONS] [--batch_size BATCH_SIZE]
[--estimate_sign_with_decon] [--ignore_mode IGNORE_MODE] [--ideal_empirical_psf IDEAL_EMPIRICAL_PSF] [--cpu_workers CPU_WORKERS] [--cluster] [--partition PARTITION] [--docker]
[--digital_rotations DIGITAL_ROTATIONS] [--psf_type PSF_TYPE] [--min_psnr MIN_PSNR] [--estimated_object_gaussian_sigma ESTIMATED_OBJECT_GAUSSIAN_SIGMA] [--denoiser DENOISER]
model input dm_calibration
positional arguments:
model path to pretrained tensorflow model
input path to input .tif file
dm_calibration path DM dm_calibration mapping matrix (eg. Zernike_Korra_Bax273.csv)
options:
-h, --help show this help message and exit
--current_dm CURRENT_DM
optional path to current DM .csv file (Default: `blank mirror`)
--prev PREV previous predictions .csv file (Default: `None`)
--lateral_voxel_size LATERAL_VOXEL_SIZE
lateral voxel size in microns for X (Default: `0.097`)
--axial_voxel_size AXIAL_VOXEL_SIZE
axial voxel size in microns for Z (Default: `0.1`)
--wavelength WAVELENGTH
wavelength in microns (Default: `0.51`)
--dm_damping_scalar DM_DAMPING_SCALAR
scale DM actuators by an arbitrary multiplier (Default: `0.75`)
--freq_strength_threshold FREQ_STRENGTH_THRESHOLD
minimum frequency threshold in fourier space (percentages; values below that will be set to the desired minimum) (Default: `0.01`)
--prediction_threshold PREDICTION_THRESHOLD
set predictions below threshold to zero (waves) (Default: `0.0`)
--confidence_threshold CONFIDENCE_THRESHOLD
optional threshold to flag unconfident predictions based on the standard deviations of the predicted amplitudes for all digital rotations (microns) (Default: `0.02`)
--sign_threshold SIGN_THRESHOLD
flip sign of modes above given threshold relative to your initial prediction (Default: `0.9`)
--plot a toggle for plotting predictions
--plot_rotations a toggle for plotting predictions for digital rotations
--num_predictions NUM_PREDICTIONS
number of predictions per sample to estimate model's confidence (Default: `1`)
--batch_size BATCH_SIZE
maximum batch size for the model (Default: `100`)
--estimate_sign_with_decon
a toggle for estimating signs of each Zernike mode via decon
--ignore_mode IGNORE_MODE
ANSI index for mode you wish to ignore (Default: `[0, 1, 2, 4]`)
--ideal_empirical_psf IDEAL_EMPIRICAL_PSF
path to an ideal empirical psf (Default: `None` ie. will be simulated automatically)
--cpu_workers CPU_WORKERS
number of CPU cores to use (Default: `-1`)
--cluster a toggle to run predictions on our cluster
--partition PARTITION
slurm partition to use on the ABC cluster (Default: `abc_a100`)
--docker a toggle to run predictions through docker container
--digital_rotations DIGITAL_ROTATIONS
optional flag for applying digital rotations (Default: `361`)
--psf_type PSF_TYPE widefield, 2photon, confocal, or a path to an LLS excitation profile (Default: None; to keep default mode used during training)
--min_psnr MIN_PSNR Will blank image if filtered image does not meet this SNR minimum. min_psnr=0 disables this threshold (Default: `5`)
--estimated_object_gaussian_sigma ESTIMATED_OBJECT_GAUSSIAN_SIGMA
size of object for creating an ideal psf (default: 0; single pixel) (Default: `0.0`)
--denoiser DENOISER path to denoiser model (Default: `None`)
To tile a large FOV and predict the wavefront of each tile, you can use the predict_tiles
function:
pytest -s -v --disable-pytest-warnings --color=yes tests/test_ao.py -k test_predict_tiles
tests/test_ao.py::test_predict_tiles Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.125_y_0.125_z_0.2_twd_simulator_False
Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.097_y_0.097_z_0.2_twd_simulator_False
Locating tiles: [288]: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 288/288 [00:41<00:00, 6.90 tile/s] 41.7s elapsed
Generate fourier embeddings: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 171/171 [00:40<00:00, 4.17 .tif file/s] 41.0s elapsed
965/965 [==============================] - 950s 976ms/step
Evaluate predictions: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 171/171 [00:00<00:00, 857925.82 evals/s] 0.0s elapsed
PASSED
You can then run aggregate_predictions
to create some visualizations :
pytest -s -v --disable-pytest-warnings --color=yes tests/test_ao.py -k test_aggregate_tiles
tests/test_ao.py::test_aggregate_tiles Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.51_na_1.0_ri_1.33_x_0.097_y_0.097_z_0.2_twd_simulator_False
Loading cached SyntheticPSF instance from /app/aovift/SyntheticPSFCache/.._lattice_YuMB_NAlattice0p35_NAAnnulusMax0p40_NAsigma0p1.mat_shape_64-64-64_lam_0.59_na_1.0_ri_1.33_x_0.097_y_0.097_z_0.2_twd_simulator_False
Number of tiles in each cluster of aggregated map, z=0
c count
1 8
2 23
3 12
4 14
5 30
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 288/288 [00:08<00:00, 34.95it/s]
PASSED
usage: ao.py predict_tiles [-h] [--current_dm CURRENT_DM] [--batch_size BATCH_SIZE] [--window_size WINDOW_SIZE] [--prev PREV] [--lateral_voxel_size LATERAL_VOXEL_SIZE]
[--axial_voxel_size AXIAL_VOXEL_SIZE] [--wavelength WAVELENGTH] [--freq_strength_threshold FREQ_STRENGTH_THRESHOLD] [--confidence_threshold CONFIDENCE_THRESHOLD]
[--sign_threshold SIGN_THRESHOLD] [--estimated_object_gaussian_sigma ESTIMATED_OBJECT_GAUSSIAN_SIGMA] [--plot] [--plot_rotations] [--num_predictions NUM_PREDICTIONS]
[--estimate_sign_with_decon] [--ignore_mode IGNORE_MODE] [--ideal_empirical_psf IDEAL_EMPIRICAL_PSF] [--cpu_workers CPU_WORKERS] [--cluster] [--partition PARTITION] [--docker]
[--digital_rotations DIGITAL_ROTATIONS] [--shift SHIFT] [--psf_type PSF_TYPE] [--min_psnr MIN_PSNR] [--denoiser DENOISER]
model input dm_calibration
positional arguments:
model path to pretrained tensorflow model
input path to input .tif file
dm_calibration path DM dm_calibration mapping matrix (eg. Zernike_Korra_Bax273.csv)
options:
-h, --help show this help message and exit
--current_dm CURRENT_DM
optional path to current DM .csv file (Default: `blank mirror`)
--batch_size BATCH_SIZE
maximum batch size for the model (Default: `100`)
--window_size WINDOW_SIZE
size of the window to crop each tile (Default: `64-64-64`)
--prev PREV previous predictions .csv file (Default: `None`)
--lateral_voxel_size LATERAL_VOXEL_SIZE
lateral voxel size in microns for X (Default: `0.097`)
--axial_voxel_size AXIAL_VOXEL_SIZE
axial voxel size in microns for Z (Default: `0.1`)
--wavelength WAVELENGTH
wavelength in microns (Default: `0.51`)
--freq_strength_threshold FREQ_STRENGTH_THRESHOLD
minimum frequency threshold in fourier space (percentages; values below that will be set to the desired minimum) (Default: `0.01`)
--confidence_threshold CONFIDENCE_THRESHOLD
optional threshold to flag unconfident predictions based on the standard deviations of the predicted amplitudes for all digital rotations (microns) (Default: `0.015`)
--sign_threshold SIGN_THRESHOLD
flip sign of modes above given threshold relative to your initial prediction (Default: `0.9`)
--estimated_object_gaussian_sigma ESTIMATED_OBJECT_GAUSSIAN_SIGMA
size of object for creating an ideal psf (default: 0; single pixel) (Default: `0.0`)
--plot a toggle for plotting predictions
--plot_rotations a toggle for plotting predictions for digital rotations
--num_predictions NUM_PREDICTIONS
number of predictions per tile to estimate model's confidence (Default: `1`)
--estimate_sign_with_decon
a toggle for estimating signs of each Zernike mode via decon
--ignore_mode IGNORE_MODE
ANSI index for mode you wish to ignore (Default: `[0, 1, 2, 4]`)
--ideal_empirical_psf IDEAL_EMPIRICAL_PSF
path to an ideal empirical psf (Default: `None` ie. will be simulated automatically)
--cpu_workers CPU_WORKERS
number of CPU cores to use (Default: `-1`)
--cluster a toggle to run predictions on our cluster
--partition PARTITION
slurm partition to use on the ABC cluster (Default: `abc_a100`)
--docker a toggle to run predictions through docker container
--digital_rotations DIGITAL_ROTATIONS
optional flag for applying digital rotations (Default: `361`)
--shift SHIFT optional flag for applying digital x shift (Default: `0`)
--psf_type PSF_TYPE widefield, 2photon, confocal, or a path to an LLS excitation profile (Default: None; to keep default mode used during training)
--min_psnr MIN_PSNR Will blank image if filtered image does not meet this SNR minimum. min_psnr=0 disables this threshold (Default: `5`)
--denoiser DENOISER path to denoiser model (Default: `None`)
usage: ao.py aggregate_predictions [-h] [--current_dm CURRENT_DM] [--dm_damping_scalar DM_DAMPING_SCALAR] [--prediction_threshold PREDICTION_THRESHOLD] [--majority_threshold MAJORITY_THRESHOLD]
[--aggregation_rule AGGREGATION_RULE] [--min_percentile MIN_PERCENTILE] [--max_percentile MAX_PERCENTILE] [--max_isoplanatic_clusters MAX_ISOPLANATIC_CLUSTERS] [--plot]
[--ignore_tile IGNORE_TILE] [--cpu_workers CPU_WORKERS] [--cluster] [--partition PARTITION] [--docker] [--psf_type PSF_TYPE]
input dm_calibration
positional arguments:
input path to csv file
dm_calibration path DM calibration mapping matrix (eg. Zernike_Korra_Bax273.csv)
options:
-h, --help show this help message and exit
--current_dm CURRENT_DM
optional path to current DM current_dm .csv file (Default: `blank mirror`)
--dm_damping_scalar DM_DAMPING_SCALAR
scale DM actuators by an arbitrary multiplier (Default: `0.75`)
--prediction_threshold PREDICTION_THRESHOLD
set predictions below threshold to zero (p2v waves) (Default: `0.25`)
--majority_threshold MAJORITY_THRESHOLD
majority rule to use to determine dominant modes among ROIs (Default: `0.5`)
--aggregation_rule AGGREGATION_RULE
rule to use to calculate final prediction [mean, median, min, max] (Default: `median`)
--min_percentile MIN_PERCENTILE
minimum percentile to filter out outliers (Default: `5`)
--max_percentile MAX_PERCENTILE
maximum percentile to filter out outliers (Default: `95`)
--max_isoplanatic_clusters MAX_ISOPLANATIC_CLUSTERS
maximum number of unique isoplanatic patchs for clustering tiles (Default: `3`)
--plot a toggle for plotting predictions
--ignore_tile IGNORE_TILE
IDs [e.g., "z0-y0-x0"] for tiles you wish to ignore
--cpu_workers CPU_WORKERS
number of CPU cores to use (Default: `-1`)
--cluster a toggle to run predictions on our cluster
--partition PARTITION
slurm partition to use on the ABC cluster (Default: `abc_a100`)
--docker a toggle to run predictions through docker container
--psf_type PSF_TYPE widefield, 2photon, confocal, or a path to an LLS excitation profile (Default: None; to keep default mode used during training)
@article{alshaabi2025fourier,
title={Fourier-Based 3D Multistage Transformer for Aberration Correction in Multicellular Specimens},
author={Thayer Alshaabi and Daniel E. Milkie and Gaoxiang Liu and Cyna Shirazinejad and Jason L. Hong and Kemal Achour and Frederik Görlitz and Ana Milunovic-Jevtic and Cat Simmons and Ibrahim S. Abuzahriyeh and Erin Hong and Samara Erin Williams and Nathanael Harrison and Evan Huang and Eun Seok Bae and Alison N. Killilea and David G. Drubin and Ian A. Swinburne and Srigokul Upadhyayula and Eric Betzig},
journal={arXiv preprint arXiv:2503.12593},
year={2025},
url={https://arxiv.org/abs/2503.12593},
}
This work is licensed under the BSD 2-Clause License