From 6ab819da1541aa2def621b722e6d2fa5e54cd68c Mon Sep 17 00:00:00 2001
From: QL-boy <46901341+QL-boy@users.noreply.github.com>
Date: Sat, 17 Jan 2026 19:27:09 +0800
Subject: [PATCH] Implement histogram matching and seamline computation

Added histogram matching and optimal seamline computation functions.
---
 utils/blending.py | 335 +++++++++++++++++++---------------------------
 1 file changed, 141 insertions(+), 194 deletions(-)

diff --git a/utils/blending.py b/utils/blending.py
index 78b27ad..fbebb75 100644
--- a/utils/blending.py
+++ b/utils/blending.py
@@ -1,217 +1,164 @@
 import torch
 import torch.nn.functional as F
-import math
+import numpy as np
 
-def generate_weight_mask(shape, blend_mode="gaussian", feather=0.25, device="cpu"):
+def match_histograms(source, reference):
     """
-    Generates a weight mask with correct feathering logic.
-    shape: (H, W) or (C, H, W)
+    Reinhard color transfer: Match color distribution of source to reference.
+    Input/Output: [H, W, C] Tensor, Range [0,1]
     """
-    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')
+    if source.dtype != torch.float32: source = source.float()
+    if reference.dtype != torch.float32: reference = reference.float()
+
+    s_mean = torch.mean(source, dim=(0, 1), keepdim=True)
+    s_std = torch.std(source, dim=(0, 1), keepdim=True) + 1e-6
+    r_mean = torch.mean(reference, dim=(0, 1), keepdim=True)
+    r_std = torch.std(reference, dim=(0, 1), keepdim=True) + 1e-6
     
-    # 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
+    s_matched = (source - s_mean) / s_std * r_std + r_mean
+    return torch.clamp(s_matched, 0.0, 1.0)
+
+def compute_optimal_seamline(overlap_left, overlap_right, direction='horizontal', search_radius=0, seam_mode='optimal'):
+    """
+    Compute optimal seamline using Graph Cut / Dynamic Programming.
+    """
+    device = overlap_left.device
     
-    # Use max dist on either axis (Linear box falloff) or Euclidean (Circular)
-    # Usually box falloff is better for rectangular tiles
+    # [H, W]
+    diff = torch.abs(overlap_left - overlap_right).mean(dim=-1)
     
-    dist_x = torch.abs(grid_x)
-    dist_y = torch.abs(grid_y)
+    energy_map = diff.cpu().numpy()
+    H, W = energy_map.shape
     
-    # 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)
+    if seam_mode == 'middle':
+        seam_mask = torch.zeros((H, W), dtype=torch.float32, device=device)
+        if direction == 'horizontal':
+            seam_mask[:, :W // 2] = 1.0
         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)
+            seam_mask[:H // 2, :] = 1.0
+        return seam_mask
 
-    # Normalize mask range 0-1
-    mask = torch.clamp(mask, 0, 1)
-
-    if len(shape) == 3:
-        return mask.unsqueeze(0)
+    # -- MODE: OPTIMAL --
+    energy_range = np.max(energy_map) - np.min(energy_map)
+    penalty_scale = 0.01 * (energy_range + 1e-6)
+    
+    if direction == 'horizontal':
+        x_indices = np.arange(W)
+        dist_map = np.abs(x_indices - W // 2)
+        energy_map += (dist_map[None, :] * penalty_scale)
     else:
-        return mask
+        y_indices = np.arange(H)
+        dist_map = np.abs(y_indices - H // 2)
+        energy_map += (dist_map[:, None] * penalty_scale)
 
-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)
+    if search_radius > 0:
+        large_val = 1e6
+        if direction == 'horizontal': 
+            center = W // 2
+            start = max(0, center - search_radius)
+            end = min(W, center + search_radius)
+            energy_map[:, :start] += large_val
+            energy_map[:, end:] += large_val
+        else:
+            center = H // 2
+            start = max(0, center - search_radius)
+            end = min(H, center + search_radius)
+            energy_map[:start, :] += large_val
+            energy_map[end:, :] += large_val
+
+    if direction == 'horizontal':
+        dp = energy_map.copy()
+        backtrack = np.zeros_like(dp, dtype=np.int32)
         
-    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))
+        for i in range(1, H):
+            for j in range(W):
+                j_start = max(0, j-1)
+                j_end = min(W, j+2)
+                prev_slice = dp[i-1, j_start:j_end]
+                min_idx_local = np.argmin(prev_slice)
+                dp[i, j] += prev_slice[min_idx_local]
+                backtrack[i, j] = j_start + min_idx_local
+        
+        mask_np = np.zeros((H, W), dtype=np.float32)
+        curr_j = np.argmin(dp[-1, :])
+        mask_np[-1, :curr_j] = 1.0
+        for i in range(H-1, -1, -1):
+            mask_np[i, :curr_j] = 1.0
+            if i > 0: curr_j = backtrack[i, curr_j]
             
-    return splits
+    else: # vertical
+        dp = energy_map.copy()
+        backtrack = np.zeros_like(dp, dtype=np.int32)
+        
+        for j in range(1, W):
+            for i in range(H):
+                i_start = max(0, i-1)
+                i_end = min(H, i+2)
+                prev_slice = dp[i_start:i_end, j-1]
+                min_idx_local = np.argmin(prev_slice)
+                dp[i, j] += prev_slice[min_idx_local]
+                backtrack[i, j] = i_start + min_idx_local
+                
+        mask_np = np.zeros((H, W), dtype=np.float32)
+        curr_i = np.argmin(dp[:, -1])
+        mask_np[:curr_i, -1] = 1.0
+        for j in range(W-1, -1, -1):
+            mask_np[:curr_i, j] = 1.0
+            if j > 0: curr_i = backtrack[curr_i, j]
 
-# --- Laplacian Pyramid Blending Utilities ---
+    return torch.from_numpy(mask_np).to(device)
 
-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"):
+def multi_band_blend(img_a, img_b, mask, levels=5):
     """
-    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)
+    Laplacian Pyramid Blending with shape safety.
     """
-    H, W = canvas_shape
-    C = tiles[0].shape[1]
+    if img_a.ndim == 3: img_a = img_a.unsqueeze(0)
+    if img_b.ndim == 3: img_b = img_b.unsqueeze(0)
+    if mask.ndim == 3: mask = mask.unsqueeze(0)
     
-    # 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
+    gauss_a = [img_a]
+    gauss_b = [img_b]
+    gauss_m = [mask]
     
-    # Create empty canvas levels
-    # Note: Dimensions change per level
-    curr_h, curr_w = H, W
+    min_dim = min(img_a.shape[2], img_a.shape[3])
+    max_levels = int(np.log2(min_dim)) - 1
+    levels = min(levels, max_levels)
+    
+    # 1. Build Gaussian Pyramid
     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 = []
+        gauss_a.append(F.avg_pool2d(gauss_a[-1], 2))
+        gauss_b.append(F.avg_pool2d(gauss_b[-1], 2))
+        gauss_m.append(F.avg_pool2d(gauss_m[-1], 2))
+    
+    # 2. Build Laplacian Pyramid
+    lap_a = []
+    lap_b = []
+    
     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)
+        # Shape of the CURRENT level
+        h, w = gauss_a[i].shape[2], gauss_a[i].shape[3]
+        # Upsample the NEXT level to match CURRENT
+        up_a = F.interpolate(gauss_a[i+1], size=(h, w), mode='bilinear', align_corners=False)
+        up_b = F.interpolate(gauss_b[i+1], size=(h, w), mode='bilinear', align_corners=False)
+        
+        # Now shapes match exactly
+        lap_a.append(gauss_a[i] - up_a)
+        lap_b.append(gauss_b[i] - up_b)
+    
+    lap_a.append(gauss_a[-1])
+    lap_b.append(gauss_b[-1])
+    
+    # 3. Blend
+    blend_lap = []
+    for i in range(levels + 1):
+        m = gauss_m[i]
+        blend_lap.append(lap_a[i] * (1.0 - m) + lap_b[i] * m)
+    
+    # 4. Reconstruct
+    recons = blend_lap[-1]
+    for i in range(levels - 1, -1, -1):
+        h, w = blend_lap[i].shape[2], blend_lap[i].shape[3]
+        recons = F.interpolate(recons, size=(h, w), mode='bilinear', align_corners=False)
+        recons = recons + blend_lap[i]
+        
+    return torch.clamp(recons, 0.0, 1.0)