Skip to content

cell-observatory/aovift

Repository files navigation

AOViFT: Adaptive Optical Vision Fourier Transformer

arXiv package python tensorflow license issues pr

Fourier-Based 3D Multistage Transformer for Aberration Correction in Multicellular Specimens

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*

1HHMI, 2UC Berkeley, 3Lawrence Berkeley National Laboratory, 4Chan Zuckerberg Biohub, 5Helen Wills Neuroscience Institute

arXiv   Docker   pretrained-models   Pytest   BibTeX

Overview

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.

Installation

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 docker image on a cluster via Apptainer

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

Clone repository to your host system

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.

Getting started

Note: To run AOViFT pytests you need to download the following files:

The [src/python ao.py](src/python ao.py) script provides a CLI for running our models on a given 3D stack (.tif file).

Running docker image

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.

Fourier embedding

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

Small FOV prediction

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`)

Tile-based prediction

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)

BibTeX

@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},
}

License

This work is licensed under the BSD 2-Clause License