| name | generative-imaging |
| description | Restore images (denoising, inpainting, super-resolution) via DDRM SVD data-consistency projection with DDIM diffusion sampling in NumPy. Use for cryo-EM/MRI restoration or inverse problems y=Hx with a diffusion prior. |
| tool_type | python |
| primary_tool | NumPy |
Generative Imaging
When to Use
- Formulating image restoration (denoising, inpainting, super-resolution) as an inverse problem
y = Hx + noise solved with a diffusion prior (DDRM)
- Sampling from a diffusion model with DDIM to reconstruct an image in fewer steps than training
- Restoring degraded scientific images — cryo-EM micrographs, MRI k-space undersampling, fluorescence microscopy — where the forward degradation operator
H is known
- Building/debugging noise schedules (linear vs cosine) or checking whether ᾱₜ indexing is inverted
- Visualizing a score field
∇ₓ log p(x) to build intuition for reverse diffusion
Version Compatibility
- Python ≥3.10, NumPy ≥1.24 (all patterns below are pure NumPy — no GPU/framework required)
- Concepts follow DDPM (Ho et al. 2020), DDIM (Song et al. 2021), and DDRM (Kawar et al. 2022)
- For a trainable neural denoiser instead of the oracle used here, pair with PyTorch (see
ai-science-diffusion-generative-models)
Prerequisites
pip install numpy (add matplotlib for the plots in the examples)
- Familiarity with Gaussian diffusion basics: forward process
xₜ = √ᾱₜ x₀ + √(1-ᾱₜ) ε
- Linear algebra: SVD, pseudoinverse, weighted least squares
Goal: turn a noisy/masked/downsampled measurement y = Hx₀ + σ·ε into a restored image x̂₀ using a diffusion prior instead of a hand-tuned regularizer.
Approach: run DDIM reverse steps with a denoiser that predicts x₀, then after each step project the prediction onto the measurement-consistent subspace via the SVD of H (DDRM), and feed the corrected x₀ back into the next reverse step.
import numpy as np
def linear_schedule(T=1000, beta_start=1e-4, beta_end=0.02):
"""Standard DDPM linear noise schedule."""
betas = np.linspace(beta_start, beta_end, T)
alphas = 1 - betas
alpha_bars = np.cumprod(alphas)
return betas, alphas, alpha_bars
def cosine_schedule(T=1000, s=0.008):
"""Improved cosine schedule (Nichol & Dhariwal 2021); smoother than linear."""
t = np.arange(T + 1)
f = np.cos((t / T + s) / (1 + s) * np.pi / 2) ** 2
alpha_bars = f / f[0]
betas = np.clip(1 - alpha_bars[1:] / alpha_bars[:-1], 0, 0.999)
return betas, 1 - betas, alpha_bars[1:]
def forward_diffuse(x0, t, alpha_bars, rng=None):
"""Add noise to x0 at timestep t: xt = sqrt(abar_t) x0 + sqrt(1-abar_t) eps."""
rng = rng or np.random.default_rng()
ab = alpha_bars[t]
eps = rng.standard_normal(x0.shape)
xt = np.sqrt(ab) * x0 + np.sqrt(1 - ab) * eps
return xt, eps
def ddim_step(xt, eps_pred, t, t_prev, alpha_bars, eta=0.0):
"""One DDIM reverse step. eta=0 -> deterministic; eta=1 -> DDPM-like stochastic."""
ab_t = alpha_bars[t]
ab_prev = alpha_bars[t_prev] t_prev >=
x0_pred = (xt - np.sqrt( - ab_t) * eps_pred) / (np.sqrt(ab_t) + )
x0_pred = np.clip(x0_pred, -, )
sigma = eta * np.sqrt(( - ab_prev) / ( - ab_t + ) * ( - ab_t / ab_prev))
noise = sigma * np.random.randn(*xt.shape) eta >
xt_prev = (np.sqrt(ab_prev) * x0_pred
+ np.sqrt(( - ab_prev - sigma**, )) * eps_pred
+ noise)
xt_prev, x0_pred
Goal: enforce that the reconstructed image stays consistent with the observed measurement y at every reverse step (DDRM data consistency).
Approach: decompose H with SVD once, then at each step blend the diffusion prior's x0_pred toward y in proportion to the singular-value SNR — well-observed directions (large s) trust y more, ill-posed directions (small/zero s) trust the prior more.
import numpy as np
def ddrm_projection(x0_pred, y_obs, H, sigma_obs=0.05):
"""
General DDRM data-consistency step for a linear degradation y = H @ x0 + sigma_obs * noise.
H: (m, n) degradation matrix (identity=denoising, row-subset=inpainting,
block-average=downsampling for super-resolution). Uses the FULL V basis so that
directions H can't observe (its null space) keep the prior's x0_pred untouched --
a thin SVD would silently drop them.
"""
U, s, Vt = np.linalg.svd(H, full_matrices=True)
x0_flat = x0_pred.flatten()
v_coeff = Vt @ x0_flat
y_coeff = U.T @ y_obs
m = len(s)
safe_s = np.where(s > 1e-8, s, 1.0)
trust = np.where(s > 1e-8, s**2 / (s**2 + sigma_obs**2), 0.0)
v_coeff_updated = v_coeff.copy()
v_coeff_updated[:m] = (1 - trust) * v_coeff[:m] + trust * (y_coeff / safe_s)
return (Vt.T @ v_coeff_updated).reshape(x0_pred.shape)
def ddrm_inpainting_step(x0_pred, y_masked, mask, sigma_obs=0.05):
"""Special case of DDRM projection for a binary mask (H = diag(mask)); avoids full SVD."""
scale = 1 / (1 + sigma_obs**2)
x0_proj = x0_pred.copy()
observed = mask == 1
x0_proj[observed] = x0_pred[observed] + scale * (y_masked[observed] - x0_pred[observed])
return x0_proj
def ddim_ddrm_sample(denoiser, xT, y_obs, mask_or_H, alpha_bars, timesteps,
is_mask=, sigma_obs=):
xt = xT.copy()
i, t (timesteps[:-]):
t_prev = timesteps[i + ]
eps_pred = denoiser(xt, t)
_, x0_pred = ddim_step(xt, eps_pred, t, t_prev, alpha_bars)
is_mask:
x0_proj = ddrm_inpainting_step(x0_pred, y_obs, mask_or_H, sigma_obs)
:
x0_proj = ddrm_projection(x0_pred, y_obs, mask_or_H, sigma_obs)
ab_prev = alpha_bars[t_prev] t_prev >=
xt = np.sqrt(ab_prev) * x0_proj + np.sqrt( - ab_prev) * eps_pred
xt
Goal: sanity-check a restoration setup end to end on a toy 1D "image" before wiring in a real network.
Approach: build ground truth, corrupt it with a known mask, run the loop above with an oracle denoiser (one that knows x0, standing in for a trained network), and confirm the reconstruction error drops versus the naive noisy observation.
def demo():
"""Self-check: DDRM inpainting should beat the naive masked observation."""
rng = np.random.default_rng(0)
n = 64
x0_true = np.sin(np.linspace(0, 4 * np.pi, n)) * 0.8
_, _, alpha_bars = cosine_schedule(T=1000)
xT, _ = forward_diffuse(x0_true, 999, alpha_bars, rng)
timesteps = np.linspace(999, 1, 20).astype(int)
mask = np.ones(n)
mask[n // 4: 3 * n // 4] = 0
y_masked = x0_true * mask
def oracle_denoiser(xt, t):
ab = alpha_bars[t]
return (xt - np.sqrt(ab) * x0_true) / np.sqrt(1 - ab + 1e-9)
x_hat = ddim_ddrm_sample(oracle_denoiser, xT, y_masked, mask, alpha_bars,
timesteps, is_mask=True)
mse_masked_region = ((x_hat[mask == 0] - x0_true[mask == 0]) ** 2).mean()
assert mse_masked_region < 0.5, f"DDRM inpainting failed to converge: MSE={mse_masked_region:.3f}"
print(f"DDRM inpainting MSE in masked region: {mse_masked_region:.4f}")
H_identity = np.eye(n)
y_noisy = x0_true + 0.3 * rng.standard_normal(n)
x0_proj = ddrm_projection(x0_true, y_noisy, H_identity, sigma_obs=)
mse_noisy = ((y_noisy - x0_true) ** ).mean()
mse_proj = ((x0_proj - x0_true) ** ).mean()
mse_proj < mse_noisy,
()
__name__ == :
demo()
Pitfalls
- Tensor shape:
betas/alpha_bars must be shape (T,), not (T, 1) — broadcasting against images fails silently
- ᾱₜ indexing convention:
alpha_bars[T-1] is most noisy, alpha_bars[0] is least noisy; mixing this up flips the whole schedule
- Clipping
x0_pred: without clipping to a sane range (e.g. [-1, 1] or [-1.5, 1.5]), errors compound over DDIM steps and the image diverges
- SVD conditioning: near-zero singular values in
H (e.g. aggressive downsampling) make the pseudoinverse unstable — regularize with sigma_obs or truncate small singular values
- Re-deriving
xt after projection: skipping the re-projection step after ddrm_projection/ddrm_inpainting_step leaves the trajectory inconsistent with the corrected x0, and later steps silently undo the correction
- Cosine vs linear schedule: cosine keeps more signal for longer near
t=0, which matters for high-frequency image detail; linear schedules over-noise images too quickly for restoration tasks
See Also
ai-science-diffusion-generative-models — full DDPM/DDIM training and sampling reference
bio-applied-deep-learning-for-biology — training the neural denoiser this skill treats as an oracle
bio-structural-biology-modern-structure-prediction — cryo-EM/structure reconstruction context for imaging restoration
ai-science-llm-finetuning — contrasting autoregressive generative paradigm