m1_4kandClaude Sonnet 4.6
a46d9b3796
- k-space (Fourier-domain) random subsampling (10% default) - ISTA reconstruction with TV (Chambolle) regularisation - Sampling patterns: variable_density (MRI-style) and uniform - Adjoint operator verified to machine epsilon - scikit-image TV denoiser with pure-NumPy Chambolle fallback Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
300 lines
10 KiB
Python
300 lines
10 KiB
Python
"""
|
||
Compressed Sensing node for ComfyUI — k-space (Fourier-domain) implementation.
|
||
|
||
Measurement:
|
||
y = Ω ⊙ F x
|
||
F : 2-D orthonormal DFT (norm='ortho') → F^H F = I
|
||
Ω : binary k-space sampling mask (sampling_ratio % of coefficients)
|
||
|
||
Reconstruction — ISTA solving:
|
||
min_x ½‖Ω ⊙ (Fx − y)‖² + λ·TV(x)
|
||
|
||
Gradient of data fidelity:
|
||
∇f(x) = F^H [ Ω ⊙ (Fx − y) ]
|
||
= IFFT2( Ω ⊙ (FFT2(x) − y) ).real
|
||
|
||
Lipschitz constant L = ‖F^H Ω F‖₂ = 1 (unitary FFT + binary mask)
|
||
→ safe step size = 1.
|
||
|
||
Incoherence:
|
||
The DFT basis is maximally incoherent with the canonical (pixel) basis.
|
||
TV exploits gradient-domain sparsity of natural images.
|
||
This satisfies the standard CS incoherence requirement.
|
||
|
||
Sampling patterns
|
||
'uniform' — uniformly random k-space subsampling
|
||
'variable_density'— 2-D Gaussian density centred at DC (MRI-style);
|
||
DC component always included.
|
||
"""
|
||
|
||
import torch
|
||
import numpy as np
|
||
|
||
try:
|
||
from skimage.restoration import denoise_tv_chambolle
|
||
_HAS_SKIMAGE = True
|
||
except ImportError:
|
||
_HAS_SKIMAGE = False
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# TV proximal operator (fallback — no scikit-image)
|
||
# ---------------------------------------------------------------------------
|
||
def _chambolle_tv_prox(u: np.ndarray, weight: float, n_iter: int = 20) -> np.ndarray:
|
||
"""Chambolle (2004) iterative TV denoising. u : [H, W, C] float64."""
|
||
p = np.zeros((*u.shape[:2], 2, u.shape[2]), dtype=np.float64)
|
||
|
||
for _ in range(n_iter):
|
||
div = np.zeros_like(u)
|
||
div[:-1] += p[:-1, :, 0]
|
||
div[1:] -= p[:-1, :, 0]
|
||
div[:, :-1] += p[:, :-1, 1]
|
||
div[:, 1:] -= p[:, :-1, 1]
|
||
|
||
d = u - weight * div
|
||
gx = np.zeros_like(u); gx[:-1] = d[1:] - d[:-1]
|
||
gy = np.zeros_like(u); gy[:, :-1] = d[:, 1:] - d[:, :-1]
|
||
|
||
g = np.stack([gx, gy], axis=2)
|
||
norm = np.sqrt((g ** 2).sum(axis=2, keepdims=True))
|
||
p = (p + 0.25 * g) / (1.0 + 0.25 * norm)
|
||
|
||
div = np.zeros_like(u)
|
||
div[:-1] += p[:-1, :, 0]; div[1:] -= p[:-1, :, 0]
|
||
div[:, :-1] += p[:, :-1, 1]; div[:, 1:] -= p[:, :-1, 1]
|
||
return u - weight * div
|
||
|
||
|
||
def _tv_prox(x: np.ndarray, weight: float) -> np.ndarray:
|
||
if _HAS_SKIMAGE:
|
||
return denoise_tv_chambolle(x, weight=weight, channel_axis=-1, max_num_iter=10)
|
||
return _chambolle_tv_prox(x, weight=weight, n_iter=15)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# k-space sampling mask
|
||
# ---------------------------------------------------------------------------
|
||
def _make_kspace_mask(
|
||
H: int,
|
||
W: int,
|
||
sampling_ratio: float,
|
||
pattern: str,
|
||
rng: np.random.RandomState,
|
||
) -> np.ndarray:
|
||
"""
|
||
Returns a boolean [H, W] mask in (unshifted) k-space layout.
|
||
|
||
pattern='uniform' — uniformly random
|
||
pattern='variable_density'— 2-D Gaussian density, DC always sampled
|
||
"""
|
||
N = H * W
|
||
n_samples = max(1, int(N * sampling_ratio))
|
||
|
||
if pattern == "variable_density":
|
||
# Build 2-D Gaussian sampling density centred at DC (k=0 → corner in numpy fft layout)
|
||
# Work in shifted layout (DC at centre) then unshift.
|
||
cy, cx = H // 2, W // 2
|
||
yi, xi = np.mgrid[0:H, 0:W]
|
||
dist_sq = ((yi - cy) / (H / 2)) ** 2 + ((xi - cx) / (W / 2)) ** 2
|
||
sigma = 0.25 # controls how tightly samples cluster near DC
|
||
prob = np.exp(-dist_sq / (2 * sigma ** 2))
|
||
prob /= prob.sum()
|
||
|
||
# Sample according to density
|
||
flat_idx = rng.choice(N, size=n_samples, replace=False, p=prob.ravel())
|
||
mask_shifted = np.zeros(N, dtype=bool)
|
||
mask_shifted[flat_idx] = True
|
||
mask_shifted = mask_shifted.reshape(H, W)
|
||
|
||
# DC (centre in shifted view) always included
|
||
mask_shifted[cy, cx] = True
|
||
|
||
# Convert from shifted layout to numpy fft layout (ifftshift)
|
||
mask = np.fft.ifftshift(mask_shifted)
|
||
else:
|
||
# Uniform random
|
||
flat_idx = rng.choice(N, size=n_samples, replace=False)
|
||
mask = np.zeros(N, dtype=bool)
|
||
mask[flat_idx] = True
|
||
mask = mask.reshape(H, W)
|
||
# Always include DC (top-left corner in numpy fft layout)
|
||
mask[0, 0] = True
|
||
|
||
return mask
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Forward / adjoint operators
|
||
# ---------------------------------------------------------------------------
|
||
def _A(x: np.ndarray, mask3: np.ndarray) -> np.ndarray:
|
||
"""A(x) = Ω ⊙ FFT2(x) → complex [H, W, C]"""
|
||
return mask3 * np.fft.fft2(x, axes=(0, 1), norm="ortho")
|
||
|
||
|
||
def _AH(y_k: np.ndarray, mask3: np.ndarray) -> np.ndarray:
|
||
"""A^H(y) = IFFT2(Ω ⊙ y).real → real [H, W, C]"""
|
||
return np.fft.ifft2(mask3 * y_k, axes=(0, 1), norm="ortho").real
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ISTA reconstruction
|
||
# ---------------------------------------------------------------------------
|
||
def _reconstruct_ista(
|
||
img: np.ndarray,
|
||
mask: np.ndarray,
|
||
num_iter: int,
|
||
tv_weight: float,
|
||
) -> np.ndarray:
|
||
"""
|
||
ISTA-based CS reconstruction from k-space measurements.
|
||
|
||
Parameters
|
||
----------
|
||
img : [H, W, C] float64 — original image (used only to compute y)
|
||
mask : [H, W] bool — k-space sampling mask
|
||
num_iter : int — ISTA iterations
|
||
tv_weight: float — TV regularisation strength λ
|
||
|
||
Returns
|
||
-------
|
||
x : [H, W, C] float64, values clipped to [0, 1]
|
||
"""
|
||
mask3 = mask[:, :, np.newaxis] # [H, W, 1] → broadcast over C
|
||
|
||
# --- Measurement: sample in k-space ---
|
||
y_k = _A(img, mask3) # complex [H, W, C]
|
||
|
||
# --- Initialise with zero-filled IFFT ---
|
||
x = _AH(y_k, mask3) # real [H, W, C]
|
||
x = np.clip(x, 0.0, 1.0)
|
||
|
||
# L = 1 (unitary FFT + binary mask) → step = 1
|
||
step = 1.0
|
||
|
||
for _ in range(num_iter):
|
||
# Gradient of data-fidelity term: ∇f = A^H ( A(x) − y_k )
|
||
grad = _AH(_A(x, mask3) - y_k, mask3)
|
||
|
||
# Gradient step
|
||
x = x - step * grad
|
||
|
||
# TV proximal step (shrinks gradient-domain coefficients)
|
||
x = _tv_prox(x, weight=tv_weight * step)
|
||
|
||
x = np.clip(x, 0.0, 1.0)
|
||
|
||
return x
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# k-space mask visualisation (fftshift so DC is at centre)
|
||
# ---------------------------------------------------------------------------
|
||
def _mask_to_image(mask: np.ndarray, C: int) -> np.ndarray:
|
||
"""Convert boolean k-space mask → visible [H, W, C] float32 image."""
|
||
shifted = np.fft.fftshift(mask).astype(np.float32)
|
||
return np.stack([shifted] * C, axis=-1)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ComfyUI Node
|
||
# ---------------------------------------------------------------------------
|
||
class CompressedSensingNode:
|
||
"""
|
||
Compressed Sensing Reconstruction (k-space / Fourier-domain)
|
||
-------------------------------------------------------------
|
||
Measures 10 % of the image's 2-D Fourier (k-space) coefficients,
|
||
then recovers the image by solving a TV-regularised least-squares
|
||
problem via ISTA.
|
||
|
||
Outputs
|
||
-------
|
||
reconstructed : recovered image
|
||
kspace_mask : visualisation of the k-space sampling pattern
|
||
(DC component at centre, white = sampled)
|
||
"""
|
||
|
||
SAMPLING_PATTERNS = ["variable_density", "uniform"]
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
},
|
||
"optional": {
|
||
"sampling_ratio": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.10,
|
||
"min": 0.01,
|
||
"max": 1.00,
|
||
"step": 0.01,
|
||
"display": "slider",
|
||
},
|
||
),
|
||
"sampling_pattern": (
|
||
cls.SAMPLING_PATTERNS,
|
||
{"default": "variable_density"},
|
||
),
|
||
"iterations": (
|
||
"INT",
|
||
{"default": 300, "min": 10, "max": 2000, "step": 10},
|
||
),
|
||
"tv_weight": (
|
||
"FLOAT",
|
||
{"default": 0.05, "min": 0.001, "max": 2.0, "step": 0.005},
|
||
),
|
||
"seed": (
|
||
"INT",
|
||
{"default": 0, "min": 0, "max": 2**31 - 1},
|
||
),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "IMAGE")
|
||
RETURN_NAMES = ("reconstructed", "kspace_mask")
|
||
FUNCTION = "process"
|
||
CATEGORY = "CompressedSensing"
|
||
|
||
def process(
|
||
self,
|
||
image: torch.Tensor,
|
||
sampling_ratio: float = 0.10,
|
||
sampling_pattern: str = "variable_density",
|
||
iterations: int = 300,
|
||
tv_weight: float = 0.05,
|
||
seed: int = 0,
|
||
):
|
||
"""
|
||
Parameters
|
||
----------
|
||
image : [B, H, W, C] float32, values in [0, 1]
|
||
sampling_ratio : fraction of k-space coefficients to observe
|
||
sampling_pattern : 'variable_density' (MRI-style) | 'uniform'
|
||
iterations : ISTA iteration count
|
||
tv_weight : TV regularisation weight λ
|
||
seed : random seed (0 = different mask every run)
|
||
"""
|
||
rng = np.random.RandomState(seed if seed != 0 else None)
|
||
device = image.device
|
||
|
||
results, masks_out = [], []
|
||
|
||
for b in range(image.shape[0]):
|
||
img_np = image[b].cpu().numpy().astype(np.float64) # [H, W, C]
|
||
H, W, C = img_np.shape
|
||
|
||
# 1. k-space sampling mask
|
||
mask = _make_kspace_mask(H, W, sampling_ratio, sampling_pattern, rng)
|
||
|
||
# 2. Measure + reconstruct
|
||
x_rec = _reconstruct_ista(img_np, mask, iterations, tv_weight)
|
||
|
||
results.append(x_rec.astype(np.float32))
|
||
masks_out.append(_mask_to_image(mask, C))
|
||
|
||
out_images = torch.from_numpy(np.stack(results)).to(device)
|
||
out_masks = torch.from_numpy(np.stack(masks_out)).to(device)
|
||
|
||
return (out_images, out_masks)
|