facok
e9b7d7481d
core/bilateral.py and core/relationships.py were referenced by nodes/anchor.py but never committed, causing ModuleNotFoundError on fresh clones.
78 lines
3.0 KiB
Python
78 lines
3.0 KiB
Python
"""Bilateral filter in LCS space for smooth color anchoring."""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def estimate_bilateral_params(c, h_len, w_len):
|
|
"""Estimate bilateral filter parameters from local color statistics.
|
|
|
|
Computes per-channel spatial std of c across the grid, takes the median
|
|
to derive sigma_color. sigma_spatial is fixed at 1.5 (5x5 kernel is small).
|
|
|
|
c: [B, L, 3] LCS coordinates
|
|
Returns: (sigma_spatial, sigma_color) floats
|
|
"""
|
|
B = c.shape[0]
|
|
grid = c.reshape(B, h_len, w_len, 3) # [B, H, W, 3]
|
|
# Per-channel std across spatial dims → [B, 3]
|
|
channel_std = grid.reshape(B, -1, 3).std(dim=1) # [B, 3]
|
|
# Median across batch and channels
|
|
median_std = float(channel_std.median())
|
|
sigma_color = max(0.05, min(3.0, 0.75 * median_std))
|
|
sigma_spatial = 1.5
|
|
return sigma_spatial, sigma_color
|
|
|
|
|
|
def bilateral_filter_lcs(c, h_len, w_len, sigma_spatial, sigma_color, kernel_radius=2):
|
|
"""Bilateral filter on [B, L, 3] LCS coordinates arranged on h_len x w_len grid.
|
|
|
|
Uses spatial distance + LCS color distance as joint weights.
|
|
kernel_radius=2 -> 5x5 neighborhood (25 lookups per patch).
|
|
Returns [B, L, 3] filtered coordinates.
|
|
"""
|
|
B = c.shape[0]
|
|
# Reshape to spatial grid
|
|
grid = c.reshape(B, h_len, w_len, 3) # [B, H, W, 3]
|
|
|
|
# Pad by kernel_radius (replicate) — pad last two spatial dims
|
|
# F.pad on [B, H, W, 3]: need to pad dims -3 and -2 (H and W)
|
|
# Permute to [B, 3, H, W] for F.pad, then back
|
|
grid_chw = grid.permute(0, 3, 1, 2) # [B, 3, H, W]
|
|
r = kernel_radius
|
|
padded = F.pad(grid_chw, (r, r, r, r), mode="replicate") # [B, 3, H+2r, W+2r]
|
|
|
|
# Precompute spatial Gaussian weights for each offset in kernel
|
|
inv_2ss = -0.5 / (sigma_spatial * sigma_spatial)
|
|
inv_2sc = -0.5 / (sigma_color * sigma_color)
|
|
|
|
# Accumulate weighted sum
|
|
weight_sum = torch.zeros(B, 1, h_len, w_len, device=c.device, dtype=c.dtype)
|
|
value_sum = torch.zeros(B, 3, h_len, w_len, device=c.device, dtype=c.dtype)
|
|
|
|
for dy in range(-r, r + 1):
|
|
for dx in range(-r, r + 1):
|
|
# Spatial weight (constant per offset)
|
|
spatial_dist_sq = float(dy * dy + dx * dx)
|
|
w_spatial = torch.exp(torch.tensor(spatial_dist_sq * inv_2ss, device=c.device, dtype=c.dtype))
|
|
|
|
# Extract neighbor values from padded grid
|
|
y_start = r + dy
|
|
x_start = r + dx
|
|
neighbor = padded[:, :, y_start:y_start + h_len, x_start:x_start + w_len] # [B, 3, H, W]
|
|
|
|
# Color distance weight (per-pixel)
|
|
diff = neighbor - grid_chw # [B, 3, H, W]
|
|
color_dist_sq = (diff * diff).sum(dim=1, keepdim=True) # [B, 1, H, W]
|
|
w_color = torch.exp(color_dist_sq * inv_2sc) # [B, 1, H, W]
|
|
|
|
w = w_spatial * w_color
|
|
weight_sum = weight_sum + w
|
|
value_sum = value_sum + w * neighbor
|
|
|
|
# Normalize
|
|
result = value_sum / weight_sum.clamp(min=1e-8) # [B, 3, H, W]
|
|
|
|
# Back to [B, L, 3]
|
|
return result.permute(0, 2, 3, 1).reshape(B, -1, 3)
|