import numpy as np import torch from scipy import ndimage def phase_stretch_transform(img_gray: np.ndarray, lpf_sigma: float = 0.21, phase_strength: float = 0.48, warp_strength: float = 12.14, morph_flag: bool = True, min_thresh: float = -0.35, max_thresh: float = 0.35) -> tuple[np.ndarray, np.ndarray]: """ Phase Stretch Transform (PST) — matches JalaliLab reference implementation. lpf_sigma is treated as FWHM (not sigma) of the Gaussian LPF, consistent with the original MATLAB/Python reference code. Returns: phase : raw PST phase output (float, -pi..pi) edges : binary or soft edge map (float 0..1) """ H, W = img_gray.shape L = 0.5 x = np.linspace(-L, L, W) y = np.linspace(-L, L, H) FX, FY = np.meshgrid(x, y) R = np.sqrt(FX ** 2 + FY ** 2) # ── Low-pass filter ─────────────────────────────────────────────────────── # lpf_sigma is FWHM; convert to Gaussian sigma: FWHM = 2*sqrt(2*ln2)*sigma sigma = lpf_sigma / (2.0 * np.sqrt(2.0 * np.log(2.0))) expo = np.fft.fftshift(np.exp(-R ** 2 / (2.0 * sigma ** 2))) # Stage 1: apply LPF → real filtered image (matches reference two-stage approach) img_f = np.fft.fft2(img_gray.astype(np.float64)) img_filtered = np.real(np.fft.ifft2(img_f * expo)) # ── PST kernel (Eq.5): K = exp(j * S * phi_w), phi_w = W*R/(1+W*R) ─────── phi_w = warp_strength * R / (1.0 + warp_strength * R) # warped phase ramp ∈ [0, 1) pst_kernel = np.exp(1j * phase_strength * phi_w) # unit-amplitude pure-phase # Stage 2: FFT of filtered image → apply PST kernel → IFFT img_filtered_f = np.fft.fftshift(np.fft.fft2(img_filtered)) img_pst = np.fft.ifft2(np.fft.fftshift(img_filtered_f * pst_kernel)) # ── Phase angle ────────────────────────────────────────────────────────── phase = np.angle(img_pst) # range: -π … π if morph_flag: edges = ((phase > max_thresh) | (phase < min_thresh)).astype(np.float64) edges = ndimage.binary_closing(edges, structure=np.ones((3, 3))).astype(np.float64) edges = ndimage.binary_opening(edges, structure=np.ones((3, 3))).astype(np.float64) else: # Soft edge map: subtract noise floor first so background → black, edges → bright abs_phase = np.abs(phase) noise_floor = np.percentile(abs_phase, 50) # background level signal = np.clip(abs_phase - noise_floor, 0.0, None) p_signal = np.percentile(signal, 99.5) edges = np.clip(signal / (p_signal + 1e-9), 0.0, 1.0) return phase, edges class PhaseStretchTransformNode: """ComfyUI node – Phase Stretch Transform edge detector.""" CATEGORY = "image/processing" RETURN_TYPES = ("IMAGE", "IMAGE") RETURN_NAMES = ("edge_image", "phase_image") FUNCTION = "run" @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "lpf_sigma": ("FLOAT", { "default": 0.21, "min": 0.01, "max": 1.0, "step": 0.01, "tooltip": "Gaussian LPF FWHM (same unit as reference code). " "Higher = more smoothing before PST." }), "phase_strength": ("FLOAT", { "default": 0.48, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "PST phase strength (S). Higher = stronger edge response." }), "warp_strength": ("FLOAT", { "default": 12.14, "min": 0.0, "max": 100.0, "step": 0.1, "tooltip": "PST warp strength (W). Controls frequency warping." }), "min_thresh": ("FLOAT", { "default": -0.35, "min": -3.15, "max": 0.0, "step": 0.001, "tooltip": "Lower phase threshold for binary edge map (morph_flag=True only)" }), "max_thresh": ("FLOAT", { "default": 0.35, "min": 0.0, "max": 3.15, "step": 0.001, "tooltip": "Upper phase threshold for binary edge map (morph_flag=True only)" }), "morph_flag": ("BOOLEAN", { "default": False, "tooltip": "True = binary edges with morphological clean-up; " "False = soft continuous edge map (|phase| normalised)" }), "output_mode": (["grayscale", "rgb_overlay"], { "default": "grayscale", "tooltip": "grayscale = white edges on black; " "rgb_overlay = edges overlaid in red on original" }), }, } def run(self, image: torch.Tensor, lpf_sigma: float, phase_strength: float, warp_strength: float, min_thresh: float, max_thresh: float, morph_flag: bool, output_mode: str): """image shape: (B, H, W, C) – ComfyUI convention, float32 0..1""" B, H, W, C = image.shape edge_frames = [] phase_frames = [] for b in range(B): frame = image[b].cpu().numpy() # (H, W, C) # Grayscale conversion (luminance weights) if C == 1: gray = frame[..., 0] else: gray = (0.2126 * frame[..., 0] + 0.7152 * frame[..., 1] + 0.0722 * frame[..., 2]) phase, edges = phase_stretch_transform( gray, lpf_sigma=lpf_sigma, phase_strength=phase_strength, warp_strength=warp_strength, morph_flag=morph_flag, min_thresh=min_thresh, max_thresh=max_thresh, ) # ── edge output ────────────────────────────────────────────────── if output_mode == "rgb_overlay": overlay = frame.copy() mask = edges > 0.5 overlay[mask] = np.array([1.0, 0.2, 0.2], dtype=np.float32) edge_out = overlay.astype(np.float32) else: edge_out = np.stack([edges] * 3, axis=-1).astype(np.float32) # ── phase_image: same noise-floor-subtracted normalization ──────── abs_phase = np.abs(phase) noise_floor = np.percentile(abs_phase, 50) signal = np.clip(abs_phase - noise_floor, 0.0, None) p_signal = np.percentile(signal, 99.5) phase_norm = np.clip(signal / (p_signal + 1e-9), 0.0, 1.0) phase_out = np.stack([phase_norm] * 3, axis=-1).astype(np.float32) edge_frames.append(torch.from_numpy(edge_out)) phase_frames.append(torch.from_numpy(phase_out)) edge_tensor = torch.stack(edge_frames, dim=0) # (B, H, W, 3) phase_tensor = torch.stack(phase_frames, dim=0) # (B, H, W, 3) return edge_tensor, phase_tensor