facok
44812b79a6
- Remove dead detect_anomalies() (superseded by adaptive variant) - Remove unused r_ref param from infer_color_from_neighbors() - Remove write-only c_ema state from anchor hook - Use math.exp() instead of torch.tensor()+torch.exp() in bilateral loop - Use in-place .add_() for accumulation in bilateral filter - Pre-normalize padded tensor once in compute_local_relationships() - Fix video VAE fallback producing duplicate vectors in calibration
118 lines
4.7 KiB
Python
118 lines
4.7 KiB
Python
"""Local color relationship analysis for drift detection and correction."""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def compute_local_relationships(c, h_len, w_len, kernel_radius=2):
|
|
"""Compute per-patch relationship vector from 5x5 neighborhood.
|
|
|
|
For each patch, cosine similarity with each of up to 24 neighbors.
|
|
Returns [B, L, N_neighbors] relationship vectors where N_neighbors = (2*r+1)^2 - 1.
|
|
"""
|
|
B = c.shape[0]
|
|
r = kernel_radius
|
|
k_size = 2 * r + 1
|
|
n_neighbors = k_size * k_size - 1 # 24 for r=2
|
|
|
|
# Reshape to spatial grid
|
|
grid = c.reshape(B, h_len, w_len, 3) # [B, H, W, 3]
|
|
|
|
# Permute to [B, 3, H, W] for padding
|
|
grid_chw = grid.permute(0, 3, 1, 2) # [B, 3, H, W]
|
|
padded = F.pad(grid_chw, (r, r, r, r), mode="replicate") # [B, 3, H+2r, W+2r]
|
|
|
|
# Center values — normalize for cosine similarity
|
|
center_norm = grid_chw / grid_chw.norm(dim=1, keepdim=True).clamp(min=1e-8)
|
|
|
|
# Pre-normalize padded tensor once (avoids per-neighbor normalization in loop)
|
|
padded_norm = padded / padded.norm(dim=1, keepdim=True).clamp(min=1e-8)
|
|
|
|
# Collect cosine similarities with each neighbor
|
|
similarities = []
|
|
for dy in range(-r, r + 1):
|
|
for dx in range(-r, r + 1):
|
|
if dy == 0 and dx == 0:
|
|
continue
|
|
y_start = r + dy
|
|
x_start = r + dx
|
|
neighbor_norm = padded_norm[:, :, y_start:y_start + h_len, x_start:x_start + w_len]
|
|
# Cosine similarity per pixel
|
|
sim = (center_norm * neighbor_norm).sum(dim=1) # [B, H, W]
|
|
similarities.append(sim)
|
|
|
|
# Stack to [B, H, W, N_neighbors] -> [B, L, N_neighbors]
|
|
rel = torch.stack(similarities, dim=-1) # [B, H, W, N_neighbors]
|
|
return rel.reshape(B, -1, n_neighbors)
|
|
|
|
|
|
def detect_anomalies_adaptive(r_current, r_reference):
|
|
"""Compare current vs reference relationships with adaptive threshold.
|
|
|
|
Uses per-batch robust outlier detection: threshold = median + 3.0 * 1.4826 * MAD.
|
|
Returns anomaly_magnitude [B, L, 1] in [0, 1].
|
|
"""
|
|
# Mean absolute difference across neighbor relationships
|
|
diff = (r_current - r_reference).abs().mean(dim=-1) # [B, L]
|
|
|
|
# Per-batch robust statistics
|
|
median = diff.median(dim=-1, keepdim=True).values # [B, 1]
|
|
mad = (diff - median).abs().median(dim=-1, keepdim=True).values # [B, 1]
|
|
threshold = median + 3.0 * 1.4826 * mad # [B, 1]
|
|
|
|
# Soft ramp above threshold, normalized to [0, 1]
|
|
anomaly = (diff - threshold).clamp(min=0.0) # [B, L]
|
|
# Normalize per-batch: max anomaly → 1.0
|
|
amax = anomaly.amax(dim=-1, keepdim=True).clamp(min=1e-8) # [B, 1]
|
|
anomaly = anomaly / amax
|
|
|
|
return anomaly.unsqueeze(-1) # [B, L, 1]
|
|
|
|
|
|
def infer_color_from_neighbors(c, anomaly_mag, h_len, w_len, kernel_radius=2):
|
|
"""For anomalous patches, infer correct color from non-anomalous neighbors.
|
|
|
|
Uses inverse-anomaly weighting: patches with low anomaly contribute more.
|
|
Returns [B, L, 3] corrected colors (blended: anomalous patches get
|
|
neighbor-inferred values, non-anomalous patches keep their original).
|
|
"""
|
|
B = c.shape[0]
|
|
r = kernel_radius
|
|
|
|
# Reshape to spatial grid
|
|
grid = c.reshape(B, h_len, w_len, 3)
|
|
anom_grid = anomaly_mag.reshape(B, h_len, w_len, 1)
|
|
|
|
# Pad both grid and anomaly
|
|
grid_chw = grid.permute(0, 3, 1, 2) # [B, 3, H, W]
|
|
anom_chw = anom_grid.permute(0, 3, 1, 2) # [B, 1, H, W]
|
|
padded_c = F.pad(grid_chw, (r, r, r, r), mode="replicate")
|
|
padded_a = F.pad(anom_chw, (r, r, r, r), mode="replicate")
|
|
|
|
# Weight neighbors by how non-anomalous they are
|
|
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):
|
|
if dy == 0 and dx == 0:
|
|
continue
|
|
y_start = r + dy
|
|
x_start = r + dx
|
|
neighbor_c = padded_c[:, :, y_start:y_start + h_len, x_start:x_start + w_len]
|
|
neighbor_a = padded_a[:, :, y_start:y_start + h_len, x_start:x_start + w_len]
|
|
|
|
# Weight: 1 - anomaly (non-anomalous neighbors get high weight)
|
|
w = (1.0 - neighbor_a).clamp(min=0.01) # [B, 1, H, W]
|
|
weight_sum = weight_sum + w
|
|
value_sum = value_sum + w * neighbor_c
|
|
|
|
# Inferred color from neighbors
|
|
inferred = value_sum / weight_sum.clamp(min=1e-8) # [B, 3, H, W]
|
|
inferred = inferred.permute(0, 2, 3, 1).reshape(B, -1, 3) # [B, L, 3]
|
|
|
|
# Blend: anomalous patches use inferred, non-anomalous keep original
|
|
# anomaly_mag is [B, L, 1], range [0, ~1]
|
|
blend = anomaly_mag.clamp(0, 1)
|
|
return c * (1.0 - blend) + inferred * blend
|