QL-boy
113d066e77
Added Laplacian pyramid fusion 新增拉普拉斯金字塔融合 Refactored weight mask feathering logic 重构权重遮罩羽化逻辑
218 lines
7.8 KiB
Python
218 lines
7.8 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import math
|
|
|
|
def generate_weight_mask(shape, blend_mode="gaussian", feather=0.25, device="cpu"):
|
|
"""
|
|
Generates a weight mask with correct feathering logic.
|
|
shape: (H, W) or (C, H, W)
|
|
"""
|
|
if len(shape) == 3:
|
|
h, w = shape[1], shape[2]
|
|
else:
|
|
h, w = shape[0], shape[1]
|
|
|
|
# Grid: -1 to 1
|
|
x = torch.linspace(-1, 1, w, device=device)
|
|
y = torch.linspace(-1, 1, h, device=device)
|
|
grid_y, grid_x = torch.meshgrid(y, x, indexing='ij')
|
|
|
|
# Feather calculation
|
|
# We want a plateau of 1.0 in the center, decaying to 0.0 at edges.
|
|
# The 'feather' parameter defines the % of the radius that is decaying.
|
|
# Distance from center d: 0 -> 1
|
|
# Effective feather region starts at 1 - feather
|
|
|
|
# Use max dist on either axis (Linear box falloff) or Euclidean (Circular)
|
|
# Usually box falloff is better for rectangular tiles
|
|
|
|
dist_x = torch.abs(grid_x)
|
|
dist_y = torch.abs(grid_y)
|
|
|
|
# Combine distances based on mode
|
|
if blend_mode == "gaussian":
|
|
# Radial Euclidean distance for smooth circular falloff
|
|
d = torch.sqrt(dist_x**2 + dist_y**2)
|
|
# Sigma derived from feather.
|
|
# If feather=0.5, we want 1.0 at d=0.5 and near 0 at d=1.0
|
|
# This is hard to map exactly to feather %.
|
|
# Let's switch to a robust sigmoid-like falloff for Gaussian emulation or use standard Gaussian.
|
|
sigma = 0.5
|
|
mask = torch.exp(-(d**2) / (2 * sigma**2))
|
|
|
|
elif blend_mode == "cosine":
|
|
# Cosine window
|
|
mask_x = torch.cos(grid_x * math.pi / 2)
|
|
mask_y = torch.cos(grid_y * math.pi / 2)
|
|
mask = mask_x * mask_y
|
|
|
|
else: # linear or default (box-like feather)
|
|
# Calculates how far we are from the edge (0=center, 1=edge)
|
|
# We want: 0 to (1-feather) -> Value 1
|
|
# (1-feather) to 1 -> Value 1 down to 0
|
|
|
|
# Max distance on either axis
|
|
d = torch.max(dist_x, dist_y)
|
|
|
|
# Threshold
|
|
start_decay = 1.0 - feather
|
|
if start_decay >= 1.0:
|
|
mask = torch.ones_like(d)
|
|
else:
|
|
# Linear ramp: (d - start) / (1 - start) -> 0 to 1 in decay zone
|
|
ramp = (d - start_decay) / (1.0 - start_decay + 1e-6)
|
|
mask = 1.0 - torch.clamp(ramp, 0, 1)
|
|
|
|
# Normalize mask range 0-1
|
|
mask = torch.clamp(mask, 0, 1)
|
|
|
|
if len(shape) == 3:
|
|
return mask.unsqueeze(0)
|
|
else:
|
|
return mask
|
|
|
|
def get_tiled_splits(full_height, full_width, tile_size, overlap):
|
|
"""
|
|
Calculates tile coordinates.
|
|
Returns: list of (y, x, h, w)
|
|
"""
|
|
stride = tile_size - overlap
|
|
splits = []
|
|
|
|
y_coords = []
|
|
current_y = 0
|
|
while current_y < full_height:
|
|
y_coords.append(current_y)
|
|
if current_y + tile_size >= full_height:
|
|
break
|
|
current_y += stride
|
|
# Fix last
|
|
if y_coords[-1] + tile_size > full_height:
|
|
y_coords[-1] = max(0, full_height - tile_size)
|
|
|
|
x_coords = []
|
|
current_x = 0
|
|
while current_x < full_width:
|
|
x_coords.append(current_x)
|
|
if current_x + tile_size >= full_width:
|
|
break
|
|
current_x += stride
|
|
# Fix last
|
|
if x_coords[-1] + tile_size > full_width:
|
|
x_coords[-1] = max(0, full_width - tile_size)
|
|
|
|
for y in y_coords:
|
|
for x in x_coords:
|
|
h = min(tile_size, full_height - y)
|
|
w = min(tile_size, full_width - x)
|
|
splits.append((y, x, h, w))
|
|
|
|
return splits
|
|
|
|
# --- Laplacian Pyramid Blending Utilities ---
|
|
|
|
def gaussian_pyramid(img, levels):
|
|
pyramid = [img]
|
|
for i in range(levels - 1):
|
|
# Blur and downsample
|
|
# Kernel: [1, 4, 6, 4, 1] / 16 approximation
|
|
# Using avg_pool for speed/memory or standard gaussian kernel
|
|
# For seamless blending, simple bilinear downsample often suffices
|
|
down = F.interpolate(pyramid[-1], scale_factor=0.5, mode='bilinear', align_corners=False)
|
|
pyramid.append(down)
|
|
return pyramid
|
|
|
|
def laplacian_pyramid(img, levels):
|
|
gaussian_pyr = gaussian_pyramid(img, levels)
|
|
laplacian_pyr = []
|
|
for i in range(levels - 1):
|
|
current = gaussian_pyr[i]
|
|
next_up = F.interpolate(gaussian_pyr[i+1], size=current.shape[2:], mode='bilinear', align_corners=False)
|
|
laplacian_pyr.append(current - next_up)
|
|
laplacian_pyr.append(gaussian_pyr[-1])
|
|
return laplacian_pyr
|
|
|
|
def reconstruct_from_laplacian(lap_pyr):
|
|
current = lap_pyr[-1]
|
|
for i in range(len(lap_pyr) - 2, -1, -1):
|
|
up = F.interpolate(current, size=lap_pyr[i].shape[2:], mode='bilinear', align_corners=False)
|
|
current = lap_pyr[i] + up
|
|
return current
|
|
|
|
def laplacian_pyramid_blending(tiles, splits, canvas_shape, feather=0.25, device="cpu"):
|
|
"""
|
|
Performs Laplacian Pyramid Blending on a set of tiles.
|
|
tiles: list of tensors [1, C, H, W]
|
|
splits: list of (y, x, h, w)
|
|
canvas_shape: (H, W)
|
|
"""
|
|
H, W = canvas_shape
|
|
C = tiles[0].shape[1]
|
|
|
|
# Determine levels based on tile size
|
|
min_dim = min(tiles[0].shape[2], tiles[0].shape[3])
|
|
levels = int(math.log2(min_dim)) - 2 # Keep base ~4-8 pixels
|
|
levels = max(1, min(levels, 5)) # Cap levels to save VRAM
|
|
|
|
# Initialize Pyramids for the Canvas
|
|
# We need to accumulate Weighted Laplacians and Weights for each level
|
|
# Format: List of Tensors (one per level)
|
|
canvas_L_pyr = [] # Numerator
|
|
weight_G_pyr = [] # Denominator
|
|
|
|
# Create empty canvas levels
|
|
# Note: Dimensions change per level
|
|
curr_h, curr_w = H, W
|
|
for i in range(levels):
|
|
canvas_L_pyr.append(torch.zeros((1, C, curr_h, curr_w), device=device))
|
|
weight_G_pyr.append(torch.zeros((1, 1, curr_h, curr_w), device=device))
|
|
curr_h = int(curr_h * 0.5)
|
|
curr_w = int(curr_w * 0.5)
|
|
|
|
# Process each tile
|
|
for idx, tile in enumerate(tiles):
|
|
y, x, h, w = splits[idx]
|
|
|
|
# 1. Generate Mask for Tile
|
|
# Use simple Linear feather for mask pyramid basis
|
|
mask = generate_weight_mask((C, h, w), blend_mode="linear", feather=feather, device=device)
|
|
mask = mask.unsqueeze(0) # [1, C, h, w] (broadcasts to C usually or 1)
|
|
if mask.shape[1] != 1: mask = mask[:, 0:1, :, :] # Ensure [1, 1, h, w]
|
|
|
|
# 2. Build Pyramids for Tile and Mask
|
|
tile_L_pyr = laplacian_pyramid(tile, levels)
|
|
mask_G_pyr = gaussian_pyramid(mask, levels) # Weight uses Gaussian pyr
|
|
|
|
# 3. Accumulate into Canvas Pyramids
|
|
for i in range(levels):
|
|
t_L = tile_L_pyr[i]
|
|
m_G = mask_G_pyr[i]
|
|
|
|
# Dimensions at this level
|
|
lvl_H, lvl_W = canvas_L_pyr[i].shape[2], canvas_L_pyr[i].shape[3]
|
|
|
|
# Map coordinates to this level
|
|
scale = 1.0 / (2**i)
|
|
ly, lx = int(y * scale), int(x * scale)
|
|
lh, lw = t_L.shape[2], t_L.shape[3]
|
|
|
|
# Safety checks for rounding
|
|
if ly + lh > lvl_H: lh = lvl_H - ly
|
|
if lx + lw > lvl_W: lw = lvl_W - lx
|
|
|
|
# Add to accumulators
|
|
# L_final = Sum(L_i * W_i) / Sum(W_i)
|
|
# Accumulate numerator
|
|
canvas_L_pyr[i][:, :, ly:ly+lh, lx:lx+lw] += t_L[:, :, :lh, :lw] * m_G[:, :, :lh, :lw]
|
|
# Accumulate denominator
|
|
weight_G_pyr[i][:, :, ly:ly+lh, lx:lx+lw] += m_G[:, :, :lh, :lw]
|
|
|
|
# Normalize each level
|
|
normalized_L_pyr = []
|
|
for i in range(levels):
|
|
norm = canvas_L_pyr[i] / (weight_G_pyr[i] + 1e-6)
|
|
normalized_L_pyr.append(norm)
|
|
|
|
# Reconstruct
|
|
return reconstruct_from_laplacian(normalized_L_pyr)
|