160 lines
6.5 KiB
Python
160 lines
6.5 KiB
Python
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 = -1.0,
|
||
max_thresh: float = 0.003) -> tuple[np.ndarray, np.ndarray]:
|
||
"""
|
||
Phase Stretch Transform (PST)
|
||
|
||
Returns:
|
||
phase_image : raw phase output (float, roughly -pi..pi)
|
||
edge_image : binary or soft edge map (float 0..1)
|
||
"""
|
||
H, W = img_gray.shape
|
||
L = 0.5
|
||
|
||
fx = np.linspace(-L, L, W)
|
||
fy = np.linspace(-L, L, H)
|
||
FX, FY = np.meshgrid(fx, fy)
|
||
R = np.sqrt(FX ** 2 + FY ** 2) # radial frequency grid
|
||
|
||
# ── Low-pass filter (Gaussian) ───────────────────────────────────────────
|
||
lpf = np.exp(-0.5 * (R / (lpf_sigma + 1e-9)) ** 2)
|
||
|
||
img_fft = np.fft.fftshift(np.fft.fft2(img_gray.astype(np.float64)))
|
||
img_filtered = img_fft * lpf
|
||
|
||
# ── PST kernel (warped phase ramp) ───────────────────────────────────────
|
||
W_val = warp_strength
|
||
S_val = phase_strength
|
||
pst_kernel = (W_val * R / (1.0 + W_val * R)) * np.exp(-1j * S_val * R / (1.0 + R))
|
||
|
||
img_pst = np.fft.ifft2(np.fft.ifftshift(img_filtered * pst_kernel))
|
||
|
||
# ── Phase angle ──────────────────────────────────────────────────────────
|
||
phase = np.angle(img_pst) # range: -π … π
|
||
|
||
if morph_flag:
|
||
edges = ((phase > max_thresh) | (phase < min_thresh)).astype(np.float64)
|
||
# small morphological clean-up
|
||
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 – normalise phase to 0..1
|
||
edges = (phase - phase.min()) / (phase.max() - phase.min() + 1e-9)
|
||
|
||
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": "Low-pass filter sigma – higher = more smoothing before PST"
|
||
}),
|
||
"phase_strength": ("FLOAT", {
|
||
"default": 0.48, "min": 0.0, "max": 3.0, "step": 0.01,
|
||
"tooltip": "PST phase strength (S)"
|
||
}),
|
||
"warp_strength": ("FLOAT", {
|
||
"default": 12.14, "min": 0.0, "max": 100.0, "step": 0.1,
|
||
"tooltip": "PST warp strength (W)"
|
||
}),
|
||
"min_thresh": ("FLOAT", {
|
||
"default": -1.0, "min": -3.15, "max": 0.0, "step": 0.001,
|
||
"tooltip": "Lower phase threshold for binary edge map"
|
||
}),
|
||
"max_thresh": ("FLOAT", {
|
||
"default": 0.003, "min": 0.0, "max": 3.15, "step": 0.001,
|
||
"tooltip": "Upper phase threshold for binary edge map"
|
||
}),
|
||
"morph_flag": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "Apply morphological clean-up (binary edges); "
|
||
"False returns normalised phase image"
|
||
}),
|
||
"output_mode": (["grayscale", "rgb_overlay"], {
|
||
"default": "grayscale",
|
||
"tooltip": "grayscale = white edges on black; "
|
||
"rgb_overlay = edges overlaid 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)
|
||
|
||
# Convert to grayscale for PST (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_rgb = np.stack([edges] * 3, axis=-1).astype(np.float32)
|
||
edge_out = edge_rgb
|
||
|
||
# ── phase output (normalised to 0..1, RGB) ───────────────────────
|
||
phase_norm = (phase - phase.min()) / (phase.max() - phase.min() + 1e-9)
|
||
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
|