Bayesian MEF

Bayesian multi-exposure image fusion (MEF) is a general-purpose algorithm for robust high dynamic range (HDR) imaging under low SNR or varying illumination — in particular for phase retrieval in coherent diffractive imaging. The paper, "Bayesian multi-exposure image fusion for robust high dynamic range ptychography", details the method and its benefits for ptychography (reproducing the results).

This small library is implemented in JAX, so the same code runs on CPU/GPU/TPU. Inputs may be NumPy arrays, Python lists or JAX arrays; results come back as NumPy arrays.

pip install bayes_mef

Optionally for GPU (CUDA 12 or 13)

pip install "bayes_mef[cuda13]"  # GPU with CUDA 13; `cuda12` flag for older GPUs

After a GPU install, check the GPU is picked up with bayes_mef check gpu.

Usage

A minimal example, simulating some data (Colab):

from bayes_mef import BayesianMEF, ConventionalMEF
from skimage.data import camera
import numpy as np

truth = camera()
background = 60
times = np.array([0.1, 1, 10])  # exposure times / flux factors
threshold = 1500                # detector limit

# overexposed Poisson data from the image-formation model
data = [np.random.poisson(t * truth + background) for t in times]
data_saturated = np.clip(data, None, threshold, dtype="float")

mef_em = BayesianMEF(data_saturated, threshold, times, background)
mef_em.run(n_iter=100)
fused_em = mef_em.fused_image.copy()

# ConventionalMEF is the paper's MLE baseline, with the same interface
mef_mle = ConventionalMEF(data_saturated, threshold, times, background)
mef_mle.mle()
fused_mle = mef_mle.fused_image.copy()

Precision

Single precision by default (accurate for censoring thresholds up to ~4096). For 16-bit detectors, switch to double precision before creating any array; a warning flags the risk otherwise:

import bayes_mef
bayes_mef.enable_x64()  # or set JAX_ENABLE_X64=1

When the exposures are not known

Omit times (and threshold) and they are estimated from the data, initialised from the non-saturated pixels (flux_init="matched"). The estimate is only a starting point, so pass update_fluxes=True with it and let EM refine it:

mef_em = BayesianMEF(data_saturated, background=background, update_fluxes=True)
mef_em.run(n_iter=200)

update_fluxes defaults to False — times are used exactly as given, estimated or not — and estimating them without it warns. On the simulation above, the unrefined estimate correlates 0.16 with the truth against 0.989 once refined. flux_init="uncensored" selects the v0.1.9 initialiser, whose ratios compress under heavy censoring.

Once the background dominates the signal, i.e., the weak, low-SNR regime the paper targets: summed counts carry almost no information about the exposures, and it is EM's iterative background handling that recovers the range; a warning is raised when the estimate comes out nearly flat. Estimated fluxes are relative, so the fused image is on a relative scale. Supply the real exposure times whenever you know them.

Fusing a 4D ptychogram dataset

For ptychography, we record multiple exposures per scan position. LaunchMEF fuses every scan position with a single vectorised program (CPU or GPU):

from bayes_mef import LaunchMEF

launch_mef = LaunchMEF(
    ptychogram_stack,    # (n_exposures, n_scans, dp_x, dp_y)
    background,          # a number, one dark frame, or one per exposure
    times=None,          # None -> estimated from the data
    threshold=None,      # None -> estimated automatically
    update_fluxes=False, # True -> EM refines the fluxes; pair with estimated times
    flux_init="matched",
)

# returns fused patterns (n_scans, dp_x, dp_y) and the flux factors
fused_ptyem_stack, em_flux_factors = launch_mef.run_em(n_iter=150)

# or the conventional MLE baseline over the whole stack (just the fused patterns)
fused_ptymle_stack = launch_mef.run_mle()

Backgrounds are taken however they were recorded, on LaunchMEF and on the single stack classes alike: a scalar, one offset per exposure (n_exposures,), a single dark frame (dp_x, dp_y), one frame per exposure (n_exposures, dp_x, dp_y), or one per image (the full stack shape). Each is given its exposure axis explicitly, so a per-exposure vector is never spread along the image columns, and a mismatched shape raises instead of broadcasting into something that quietly means the wrong thing.

The old n_cpus argument is accepted but ignored (it warns), since there are no worker processes to size any more.

Scans are fused in chunks sized to the device's free memory, so a stack larger than device memory still works. Set batch_size yourself if you hit a MemoryError or fuse alongside other work (results do not depend on it):

fused, fluxes = launch_mef.run_em(n_iter, batch_size=8)

See synthetic_mef.py for detailed usage on synthetic ptychography data, and benchmarks/FINDINGS.md for a study of when each method helps and for performance benchmarks.

Reproducing results

To reproduce the ptychographic reconstructions from the paper:

1. Clone the repo:

   git clone https://github.com/microscopic-image-analysis/bayes-mef.git
   cd bayes-mef
   

2. Install the pinned dependencies with uv, then prefix commands with uv run (e.g. uv run python scripts/synthetic_mef.py):

   uv sync --locked --group scripts
   

3. Download the data from Zenodo:

   ./download_data.sh
   

4. Optional: install cupy for faster GPU reconstructions.
5. Run files from scripts/ to plot the results.

Citation

If this algorithm or the publication was useful, please cite:

@article{Kodgirwar:24,
author = {Shantanu Kodgirwar and Lars Loetgering and Chang Liu and Aleena Joseph and Leona Licht and Daniel S. Penagos Molina and Wilhelm Eschen and Jan Rothhardt and Michael Habeck},
journal = {Opt. Express},
number = {16},
pages = {28090--28099},
publisher = {Optica Publishing Group},
title = {Bayesian multi-exposure image fusion for robust high dynamic range ptychography},
volume = {32},
month = {Jul},
year = {2024},
url = {https://opg.optica.org/oe/abstract.cfm?URI=oe-32-16-28090},
doi = {10.1364/OE.524284},
}