5.1 KiB
ComfyUI Compressed Sensing Node
A custom ComfyUI node that reconstructs an image from only 10% randomly sampled k-space (Fourier) coefficients using rigorous Compressed Sensing theory.
Algorithm
1. Measurement — k-space subsampling
The image is transformed into the 2-D Fourier (k-space) domain, and only 10% of the complex coefficients are retained:
y = Ω ⊙ Fx
| Symbol | Meaning |
|---|---|
x |
Original image |
F |
2-D orthonormal DFT (numpy.fft.fft2, norm='ortho') |
Ω |
Binary k-space sampling mask (10% ones) |
y |
Observed k-space measurements |
Using the Fourier domain satisfies the incoherence condition central to CS theory: the DFT basis is maximally incoherent with the canonical (pixel) basis, and natural images have sparse gradients (TV sparsity), giving theoretical recovery guarantees.
2. Reconstruction — ISTA with TV regularisation
The image is recovered by solving the convex optimisation problem:
min_x ½‖Ω ⊙ (Fx − y)‖² + λ·TV(x)
Solved iteratively via ISTA (Iterative Shrinkage-Thresholding Algorithm):
∇f(x) = F^H [ Ω ⊙ (Fx − y) ] ← gradient of data-fidelity
= IFFT2( Ω ⊙ (FFT2(x) − y) ).real
x ← prox_{λ·TV}( x − ∇f(x) ) ← TV proximal step (Chambolle)
The step size is exactly 1 because F is unitary (F^H F = I) and Ω is a binary mask, so the Lipschitz constant of the forward operator is L = 1.
Why k-space, not pixel-domain?
| Property | Pixel-domain sampling | k-space sampling (this node) |
|---|---|---|
| Measurement domain | Pixel space | Fourier space |
| Sparsity domain | Gradient (TV) | Gradient (TV) |
| Incoherence | Weak (same domain) | Strong (maximally incoherent) |
| CS theory guarantee | Limited | Satisfies RIP w.h.p. |
| Real-world analogue | — | MRI, radar |
Sampling Patterns
variable_density (default, MRI-style)
Samples more densely near DC (low frequencies) using a 2-D Gaussian density. DC is always included. Gives better perceptual quality because most image energy is at low frequencies.
uniform
Uniformly random k-space subsampling. Theoretically cleaner (i.i.d. draws), useful for benchmarking.
Installation
# 1. Copy this folder into ComfyUI's custom_nodes directory
cp -r ComfyUI_CompressedSensing /path/to/ComfyUI/custom_nodes/
# 2. Install dependencies
pip install scikit-image
# 3. Restart ComfyUI
numpyandtorchare already bundled with ComfyUI.scikit-imageis strongly recommended for faster TV denoising. A pure-NumPy Chambolle fallback is included if it is unavailable.
Node Reference
Category: CompressedSensing
Node name: Compressed Sensing (10% Random Sampling)
Inputs
| Name | Type | Default | Description |
|---|---|---|---|
image |
IMAGE | — | Input image |
sampling_ratio |
FLOAT | 0.10 | Fraction of k-space to observe (0.01 – 1.00) |
sampling_pattern |
ENUM | variable_density |
variable_density or uniform |
iterations |
INT | 300 | ISTA iteration count. More = better quality, slower |
tv_weight |
FLOAT | 0.05 | TV regularisation strength λ. Lower = sharper; higher = smoother |
seed |
INT | 0 | Random seed. 0 = different mask every run |
Outputs
| Name | Type | Description |
|---|---|---|
reconstructed |
IMAGE | CS-reconstructed image |
kspace_mask |
IMAGE | k-space sampling pattern (DC at centre, white = sampled) |
Parameter Tuning Guide
| Goal | Suggestion |
|---|---|
| Better quality | Increase iterations (500 – 1000) |
| Sharper edges | Decrease tv_weight (0.01 – 0.03) |
| Smoother result | Increase tv_weight (0.1 – 0.3) |
| Reproducible mask | Set seed to any fixed non-zero integer |
| More measurements | Increase sampling_ratio (0.20 – 0.30) |
Example Workflow
Load Image → CompressedSensing → Preview Image (reconstructed)
→ Preview Image (kspace_mask)
Requirements
- Python 3.8+
- numpy
- torch (provided by ComfyUI)
- scikit-image ≥ 0.19 (recommended)
Theoretical Background
Compressed Sensing (Candès, Romberg & Tao 2006; Donoho 2006) states that a signal with a sparse representation in basis Ψ can be exactly recovered from m = O(s log n) incoherent measurements — far fewer than the n samples Nyquist requires — where s is the sparsity level.
This node's setup:
| CS ingredient | Choice |
|---|---|
| Signal | Natural image x ∈ ℝⁿ |
| Sparsifying basis Ψ | Gradient domain (TV) |
| Measurement matrix Φ | Subsampled 2-D DFT |
| Incoherence μ(Φ, Ψ) | Near-minimal (Fourier ↔ pixel) |
| Recovery algorithm | ISTA (convex, globally convergent) |
Lustig, Donoho & Pauly (2007) "Sparse MRI" is the seminal paper applying exactly this approach to accelerate MRI acquisition.
License
MIT