import torch import numpy as np from PIL import Image import matplotlib.pyplot as plt import matplotlib.figure as figure from matplotlib.figure import Figure import numpy.typing as npt import os import sys import tempfile import time class RegionColorMatcher: def __init__(self, factor=1.0, preserve_colors=True, preserve_luminance=True, method="adain"): """ Initialize the RegionColorMatcher. Args: factor: Strength of the color matching (0.0 to 1.0) preserve_colors: If True, convert to YUV and preserve color relationships preserve_luminance: If True, preserve the luminance when in YUV mode method: The color matching method to use (adain, mkl, hm, reinhard, mvgd, hm-mvgd-hm, hm-mkl-hm) """ self.factor = factor self.preserve_colors = preserve_colors self.preserve_luminance = preserve_luminance self.method = method def match_regions(self, img1_path, img2_path, masks1, masks2): """ Match colors between corresponding masked regions of two images. Args: img1_path: Path to first image img2_path: Path to second image masks1: Dictionary of masks for first image {label: binary_mask} masks2: Dictionary of masks for second image {label: binary_mask} Returns: A PIL Image with the color-matched result """ print(f"🎨 Color matching with method: {self.method}") print(f"📊 Processing {len(masks1)} regions from img1 and {len(masks2)} regions from img2") # Load images img1 = Image.open(img1_path).convert("RGB") img2 = Image.open(img2_path).convert("RGB") # Convert to numpy arrays and normalize to [0,1] img1_np = np.array(img1).astype(np.float32) / 255.0 img2_np = np.array(img2).astype(np.float32) / 255.0 # Create a copy of the second image as our base for color matching # We want to make img2 look like img1's colors result_np = np.copy(img2_np) # Convert images to PyTorch tensors img1_tensor = torch.from_numpy(img1_np) img2_tensor = torch.from_numpy(img2_np) result_tensor = torch.from_numpy(result_np) # Track coverage to ensure all regions are processed total_coverage = np.zeros(img2_np.shape[:2], dtype=np.float32) processed_regions = 0 # Process each mask region for label, mask1 in masks1.items(): if label not in masks2: print(f"⚠️ Skipping {label} - not found in masks2") continue mask2 = masks2[label] # Resize masks to match image dimensions if needed if mask1.shape != img1_np.shape[:2]: mask1 = self._resize_mask(mask1, img1_np.shape[:2]) if mask2.shape != img2_np.shape[:2]: mask2 = self._resize_mask(mask2, img2_np.shape[:2]) # Check mask coverage mask1_pixels = np.sum(mask1 > 0) mask2_pixels = np.sum(mask2 > 0) print(f"🔍 Processing {label}: {mask1_pixels} pixels (img1) → {mask2_pixels} pixels (img2)") if mask1_pixels == 0 or mask2_pixels == 0: print(f"⚠️ Skipping {label} - no pixels in mask") continue # Track coverage total_coverage += (mask2 > 0).astype(np.float32) processed_regions += 1 # Convert masks to torch tensors mask1_tensor = torch.from_numpy(mask1.astype(np.float32)) mask2_tensor = torch.from_numpy(mask2.astype(np.float32)) # Apply color matching for this region based on selected method if self.method == "adain": result_tensor = self._apply_adain_to_region( img1_tensor, img2_tensor, result_tensor, mask1_tensor, mask2_tensor ) else: result_tensor = self._apply_color_matcher_to_region( img1_tensor, img2_tensor, result_tensor, mask1_tensor, mask2_tensor, self.method ) print(f"✅ Completed color matching for {label}") # Debug coverage total_pixels = img2_np.shape[0] * img2_np.shape[1] covered_pixels = np.sum(total_coverage > 0) overlap_pixels = np.sum(total_coverage > 1) print(f"📊 Coverage summary:") print(f" Total image pixels: {total_pixels}") print(f" Covered pixels: {covered_pixels} ({100*covered_pixels/total_pixels:.1f}%)") print(f" Overlapping pixels: {overlap_pixels} ({100*overlap_pixels/total_pixels:.1f}%)") print(f" Processed regions: {processed_regions}") # Convert back to numpy, scale to [0,255] and convert to uint8 result_np = (result_tensor.numpy() * 255.0).astype(np.uint8) # Convert to PIL Image result_img = Image.fromarray(result_np) return result_img def _resize_mask(self, mask, target_shape): """ Resize a mask to match the target shape. Args: mask: Binary mask array target_shape: Target shape (height, width) Returns: Resized mask array """ # Convert to PIL Image for resizing mask_img = Image.fromarray((mask * 255).astype(np.uint8)) # Resize to target shape mask_img = mask_img.resize((target_shape[1], target_shape[0]), Image.NEAREST) # Convert back to numpy array and normalize to [0,1] resized_mask = np.array(mask_img).astype(np.float32) / 255.0 return resized_mask def _apply_adain_to_region(self, source_img, target_img, result_img, source_mask, target_mask): """ Apply AdaIN to match the statistics of the masked region in source to the target. Args: source_img: Source image tensor [H,W,3] (reference for color matching) target_img: Target image tensor [H,W,3] (to be color matched) result_img: Result image tensor to modify [H,W,3] source_mask: Binary mask for source image [H,W] target_mask: Binary mask for target image [H,W] Returns: Modified result tensor """ # Ensure masks are binary source_mask_binary = (source_mask > 0.5).float() target_mask_binary = (target_mask > 0.5).float() # If preserving colors, convert to YUV if self.preserve_colors: # RGB to YUV conversion matrix rgb_to_yuv = torch.tensor([ [0.299, 0.587, 0.114], [-0.14713, -0.28886, 0.436], [0.615, -0.51499, -0.10001] ]) # Convert to YUV source_yuv = torch.matmul(source_img, rgb_to_yuv.T) target_yuv = torch.matmul(target_img, rgb_to_yuv.T) result_yuv = torch.matmul(result_img, rgb_to_yuv.T) # Only normalize Y channel if preserving luminance is False channels_to_process = [0] if not self.preserve_luminance else [] # Always process U and V channels (chroma) channels_to_process.extend([1, 2]) # Process selected channels for c in channels_to_process: # Apply the color matching only to the masked region in the result result_channel = result_yuv[:,:,c] matched_channel = self._match_channel_statistics( source_yuv[:,:,c], target_yuv[:,:,c], result_channel, source_mask_binary, target_mask_binary ) # Only update the masked region in the result mask_expanded = target_mask_binary.unsqueeze(-1).expand_as(result_yuv)[:,:,c] result_yuv[:,:,c] = torch.where( mask_expanded > 0.5, matched_channel, result_channel ) # Convert back to RGB yuv_to_rgb = torch.tensor([ [1.0, 0.0, 1.13983], [1.0, -0.39465, -0.58060], [1.0, 2.03211, 0.0] ]) result_rgb = torch.matmul(result_yuv, yuv_to_rgb.T) # Only update the masked region in the result mask_expanded = target_mask_binary.unsqueeze(-1).expand_as(result_img) result_img = torch.where( mask_expanded > 0.5, result_rgb, result_img ) else: # Process each RGB channel separately for c in range(3): # Apply the color matching only to the masked region in the result result_channel = result_img[:,:,c] matched_channel = self._match_channel_statistics( source_img[:,:,c], target_img[:,:,c], result_channel, source_mask_binary, target_mask_binary ) # Only update the masked region in the result mask_expanded = target_mask_binary.unsqueeze(-1).expand_as(result_img)[:,:,c] result_img[:,:,c] = torch.where( mask_expanded > 0.5, matched_channel, result_channel ) # Ensure values are in valid range [0, 1] return torch.clamp(result_img, 0.0, 1.0) def _apply_color_matcher_to_region(self, source_img, target_img, result_img, source_mask, target_mask, method): """ Apply color-matcher library methods to match the statistics of the masked region in source to the target. Args: source_img: Source image tensor [H,W,3] (reference for color matching) target_img: Target image tensor [H,W,3] (to be color matched) result_img: Result image tensor to modify [H,W,3] source_mask: Binary mask for source image [H,W] target_mask: Binary mask for target image [H,W] method: The color matching method to use (mkl, hm, reinhard, mvgd, hm-mvgd-hm, hm-mkl-hm) Returns: Modified result tensor """ # Ensure masks are binary source_mask_binary = (source_mask > 0.5).float() target_mask_binary = (target_mask > 0.5).float() # Convert tensors to numpy arrays source_np = source_img.detach().cpu().numpy() target_np = target_img.detach().cpu().numpy() source_mask_np = source_mask_binary.detach().cpu().numpy() target_mask_np = target_mask_binary.detach().cpu().numpy() try: # Try to import the color_matcher library try: from color_matcher import ColorMatcher from color_matcher.normalizer import Normalizer except ImportError: self._install_package("color-matcher") from color_matcher import ColorMatcher from color_matcher.normalizer import Normalizer # Extract only the masked pixels from both images source_coords = np.where(source_mask_np > 0.5) target_coords = np.where(target_mask_np > 0.5) if len(source_coords[0]) == 0 or len(target_coords[0]) == 0: return result_img # Extract pixel values from masked regions source_pixels = source_np[source_coords] target_pixels = target_np[target_coords] # Initialize color matcher cm = ColorMatcher() if method == "mkl": # For MKL, calculate mean and standard deviation from masked regions source_mean = np.mean(source_pixels, axis=0) source_std = np.std(source_pixels, axis=0) target_mean = np.mean(target_pixels, axis=0) target_std = np.std(target_pixels, axis=0) # Apply the transformation result_np = np.copy(target_np) for c in range(3): # Normalize the target channel and scale by source statistics normalized = (target_np[:,:,c] - target_mean[c]) / (target_std[c] + 1e-8) * source_std[c] + source_mean[c] # Only apply to masked region result_np[:,:,c] = np.where(target_mask_np > 0.5, normalized, target_np[:,:,c]) # Convert back to tensor result_tensor = torch.from_numpy(result_np).to(result_img.device) # Blend with original based on factor result_img = torch.lerp(result_img, result_tensor, self.factor) elif method == "reinhard": # Similar to MKL but with a different normalization approach source_mean = np.mean(source_pixels, axis=0) source_std = np.std(source_pixels, axis=0) target_mean = np.mean(target_pixels, axis=0) target_std = np.std(target_pixels, axis=0) # Apply the transformation result_np = np.copy(target_np) for c in range(3): # Normalize the target channel and scale by source statistics normalized = (target_np[:,:,c] - target_mean[c]) / (target_std[c] + 1e-8) * source_std[c] + source_mean[c] # Only apply to masked region result_np[:,:,c] = np.where(target_mask_np > 0.5, normalized, target_np[:,:,c]) # Convert back to tensor result_tensor = torch.from_numpy(result_np).to(result_img.device) # Blend with original based on factor result_img = torch.lerp(result_img, result_tensor, self.factor) elif method == "mvgd": # For MVGD, we need mean and covariance matrices source_mean = np.mean(source_pixels, axis=0) source_cov = np.cov(source_pixels, rowvar=False) target_mean = np.mean(target_pixels, axis=0) target_cov = np.cov(target_pixels, rowvar=False) # Check if covariance matrices are valid if np.isnan(source_cov).any() or np.isnan(target_cov).any(): # Fallback to simple statistics matching source_std = np.std(source_pixels, axis=0) target_std = np.std(target_pixels, axis=0) result_np = np.copy(target_np) for c in range(3): normalized = (target_np[:,:,c] - target_mean[c]) / (target_std[c] + 1e-8) * source_std[c] + source_mean[c] result_np[:,:,c] = np.where(target_mask_np > 0.5, normalized, target_np[:,:,c]) else: # Apply full MVGD transformation to masked pixels # Reshape the masked pixels for matrix operations target_flat = target_np.reshape(-1, 3) result_np = np.copy(target_np) try: # Try to compute the full MVGD transformation source_cov_sqrt = np.linalg.cholesky(source_cov) target_cov_sqrt = np.linalg.cholesky(target_cov) target_cov_sqrt_inv = np.linalg.inv(target_cov_sqrt) # Compute the transformation matrix temp = target_cov_sqrt_inv @ source_cov @ target_cov_sqrt_inv.T temp_sqrt_inv = np.linalg.inv(np.linalg.cholesky(temp)) A = target_cov_sqrt @ temp_sqrt_inv @ target_cov_sqrt_inv # Apply the transformation to all pixels for i in range(target_np.shape[0]): for j in range(target_np.shape[1]): if target_mask_np[i, j] > 0.5: # Only apply to masked pixels pixel = target_np[i, j] centered = pixel - target_mean transformed = centered @ A.T + source_mean result_np[i, j] = transformed except np.linalg.LinAlgError: # Fallback to simple statistics matching source_std = np.std(source_pixels, axis=0) target_std = np.std(target_pixels, axis=0) for c in range(3): normalized = (target_np[:,:,c] - target_mean[c]) / (target_std[c] + 1e-8) * source_std[c] + source_mean[c] result_np[:,:,c] = np.where(target_mask_np > 0.5, normalized, target_np[:,:,c]) # Convert back to tensor result_tensor = torch.from_numpy(result_np).to(result_img.device) # Blend with original based on factor result_img = torch.lerp(result_img, result_tensor, self.factor) elif method in ["hm", "hm-mvgd-hm", "hm-mkl-hm"]: # For histogram-based methods, we'll create temporary cropped images with just the masked regions # Get the bounding box of the masked regions source_min_y, source_min_x = np.min(source_coords[0]), np.min(source_coords[1]) source_max_y, source_max_x = np.max(source_coords[0]), np.max(source_coords[1]) target_min_y, target_min_x = np.min(target_coords[0]), np.min(target_coords[1]) target_max_y, target_max_x = np.max(target_coords[0]), np.max(target_coords[1]) # Create cropped images with just the masked regions source_crop = source_np[source_min_y:source_max_y+1, source_min_x:source_max_x+1].copy() target_crop = target_np[target_min_y:target_max_y+1, target_min_x:target_max_x+1].copy() # Create cropped masks source_mask_crop = source_mask_np[source_min_y:source_max_y+1, source_min_x:source_max_x+1] target_mask_crop = target_mask_np[target_min_y:target_max_y+1, target_min_x:target_max_x+1] # Apply the mask to the cropped images # For non-masked areas, use the average color source_avg_color = np.mean(source_pixels, axis=0) target_avg_color = np.mean(target_pixels, axis=0) for c in range(3): source_crop[:, :, c] = np.where(source_mask_crop > 0.5, source_crop[:, :, c], source_avg_color[c]) target_crop[:, :, c] = np.where(target_mask_crop > 0.5, target_crop[:, :, c], target_avg_color[c]) try: # Use the color matcher directly on the masked regions matched_crop = cm.transfer(src=target_crop, ref=source_crop, method=method) # Apply the matched colors back to the original image, only in the masked region result_np = np.copy(target_np) # Create a mapping from crop coordinates to original image coordinates for i in range(target_crop.shape[0]): for j in range(target_crop.shape[1]): orig_i = target_min_y + i orig_j = target_min_x + j if orig_i < target_np.shape[0] and orig_j < target_np.shape[1] and target_mask_np[orig_i, orig_j] > 0.5: result_np[orig_i, orig_j] = matched_crop[i, j] # Convert back to tensor result_tensor = torch.from_numpy(result_np).to(result_img.device) # Blend with original based on factor result_img = torch.lerp(result_img, result_tensor, self.factor) except Exception as e: # Fallback to AdaIN if color matcher fails print(f"Color matcher failed for {method}, using fallback: {str(e)}") result_img = self._apply_adain_to_region( source_img, target_img, result_img, source_mask_binary, target_mask_binary ) elif method == "coral": # For CORAL method, extract masked regions and apply CORAL color transfer try: # Create masked versions of the images source_masked = source_np.copy() target_masked = target_np.copy() # Apply masks - set non-masked areas to average color source_avg_color = np.mean(source_pixels, axis=0) target_avg_color = np.mean(target_pixels, axis=0) for c in range(3): source_masked[:, :, c] = np.where(source_mask_np > 0.5, source_masked[:, :, c], source_avg_color[c]) target_masked[:, :, c] = np.where(target_mask_np > 0.5, target_masked[:, :, c], target_avg_color[c]) # Convert to torch tensors and rearrange to [C, H, W] source_tensor = torch.from_numpy(source_masked).permute(2, 0, 1).float() target_tensor = torch.from_numpy(target_masked).permute(2, 0, 1).float() # Apply CORAL color transfer matched_tensor = coral(target_tensor, source_tensor) # target gets matched to source # Convert back to [H, W, C] format matched_np = matched_tensor.permute(1, 2, 0).numpy() # Apply the matched colors back to the original image, only in the masked region result_np = np.copy(target_np) for c in range(3): result_np[:, :, c] = np.where(target_mask_np > 0.5, matched_np[:, :, c], target_np[:, :, c]) # Convert back to tensor result_tensor = torch.from_numpy(result_np).to(result_img.device) # Blend with original based on factor result_img = torch.lerp(result_img, result_tensor, self.factor) except Exception as e: # Fallback to AdaIN if CORAL fails print(f"CORAL failed for {method}, using fallback: {str(e)}") result_img = self._apply_adain_to_region( source_img, target_img, result_img, source_mask_binary, target_mask_binary ) else: # Default to AdaIN for unsupported methods result_img = self._apply_adain_to_region( source_img, target_img, result_img, source_mask_binary, target_mask_binary ) except Exception as e: # If all fails, fallback to AdaIN print(f"Error in color matching: {str(e)}, using AdaIN as fallback") result_img = self._apply_adain_to_region( source_img, target_img, result_img, source_mask_binary, target_mask_binary ) return torch.clamp(result_img, 0.0, 1.0) def _match_channel_statistics(self, source_channel, target_channel, result_channel, source_mask, target_mask): """ Match the statistics of a single channel. Args: source_channel: Source channel [H,W] (reference for color matching) target_channel: Target channel [H,W] (to be color matched) result_channel: Result channel to modify [H,W] source_mask: Binary mask for source [H,W] target_mask: Binary mask for target [H,W] Returns: Modified result channel """ # Count non-zero elements in masks source_count = torch.sum(source_mask) target_count = torch.sum(target_mask) if source_count > 0 and target_count > 0: # Calculate statistics only from masked regions source_masked = source_channel * source_mask target_masked = target_channel * target_mask # Calculate mean source_mean = torch.sum(source_masked) / source_count target_mean = torch.sum(target_masked) / target_count # Calculate variance source_var = torch.sum(((source_channel - source_mean) * source_mask) ** 2) / source_count target_var = torch.sum(((target_channel - target_mean) * target_mask) ** 2) / target_count # Calculate std (add small epsilon to avoid division by zero) source_std = torch.sqrt(source_var + 1e-8) target_std = torch.sqrt(target_var + 1e-8) # Apply AdaIN to the masked region normalized = ((target_channel - target_mean) / target_std) * source_std + source_mean # Blend with original based on factor result = torch.lerp(target_channel, normalized, self.factor) return result return result_channel def _install_package(self, package_name): """Install a package using pip.""" import subprocess subprocess.check_call([sys.executable, "-m", "pip", "install", package_name]) def create_comparison_figure(original_img, matched_img, title="Color Matching Comparison"): """ Create a matplotlib figure with the original and color-matched images. Args: original_img: Original PIL Image matched_img: Color-matched PIL Image title: Title for the figure Returns: matplotlib Figure """ fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5)) ax1.imshow(original_img) ax1.set_title("Original") ax1.axis('off') ax2.imshow(matched_img) ax2.set_title("Color Matched") ax2.axis('off') plt.suptitle(title) plt.tight_layout() return fig def coral(source, target): """ CORAL (Color Transfer using Correlated Color Temperature) implementation. Based on the original ColorMatchImage approach. Args: source: Source image tensor [C, H, W] (to be color matched) target: Target image tensor [C, H, W] (reference for color matching) Returns: Color-matched source image tensor [C, H, W] """ # Ensure tensors are float source = source.float() target = target.float() # Reshape to [C, N] where N is number of pixels C, H, W = source.shape source_flat = source.view(C, -1) # [C, H*W] target_flat = target.view(C, -1) # [C, H*W] # Compute means source_mean = torch.mean(source_flat, dim=1, keepdim=True) # [C, 1] target_mean = torch.mean(target_flat, dim=1, keepdim=True) # [C, 1] # Center the data source_centered = source_flat - source_mean # [C, H*W] target_centered = target_flat - target_mean # [C, H*W] # Compute covariance matrices N = source_centered.shape[1] source_cov = torch.mm(source_centered, source_centered.t()) / (N - 1) # [C, C] target_cov = torch.mm(target_centered, target_centered.t()) / (N - 1) # [C, C] # Add small epsilon to diagonal for numerical stability eps = 1e-5 source_cov += eps * torch.eye(C, device=source.device) target_cov += eps * torch.eye(C, device=source.device) try: # Compute the transformation matrix using Cholesky decomposition # This is more stable than eigendecomposition for positive definite matrices # Cholesky decomposition: A = L * L^T source_chol = torch.linalg.cholesky(source_cov) # Lower triangular target_chol = torch.linalg.cholesky(target_cov) # Lower triangular # Compute the transformation matrix # We want to transform source covariance to target covariance # Transform = target_chol * source_chol^(-1) source_chol_inv = torch.linalg.inv(source_chol) transform_matrix = torch.mm(target_chol, source_chol_inv) # Apply transformation: result = transform_matrix * (source - source_mean) + target_mean result_centered = torch.mm(transform_matrix, source_centered) result_flat = result_centered + target_mean # Reshape back to original shape result = result_flat.view(C, H, W) # Clamp to valid range result = torch.clamp(result, 0.0, 1.0) return result except Exception as e: # Fallback to simple mean/std matching if Cholesky fails print(f"CORAL Cholesky failed, using simple statistics matching: {e}") # Simple per-channel statistics matching source_std = torch.std(source_centered, dim=1, keepdim=True) target_std = torch.std(target_centered, dim=1, keepdim=True) # Avoid division by zero source_std = torch.clamp(source_std, min=eps) # Apply simple transformation: (source - source_mean) / source_std * target_std + target_mean result_flat = (source_centered / source_std) * target_std + target_mean result = result_flat.view(C, H, W) # Clamp to valid range result = torch.clamp(result, 0.0, 1.0) return result