""" 自动红偏校正。 纯 torch 实现的确定性白平衡：把图像转到 CIE Lab，自动挑选高亮低饱和区域 当作灰卡，测出红轴（a 通道）偏移量后只做减法校正。不调用模型或网络 API， CPU / GPU 都能跑。支持单张（连接图像端口）和批量（填写文件夹路径）两种模式。 """ from __future__ import annotations import os from typing import List, Optional import torch from PIL import Image from ..utils.image_utils import pil_to_tensor D65_WHITE = (0.95047, 1.0, 1.08883) LAB_EPSILON = 216.0 / 24389.0 LAB_KAPPA = 24389.0 / 27.0 _IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tiff", ".tif") def _load_folder_images(folder: str) -> List[Image.Image]: """从文件夹加载所有图片，返回 PIL Image 列表（RGB）。""" if not os.path.isdir(folder): raise ValueError(f"自动红偏校正：路径不是有效的文件夹：{folder}") names = sorted( n for n in os.listdir(folder) if n.lower().endswith(_IMAGE_EXTS) and os.path.isfile(os.path.join(folder, n)) ) if not names: raise ValueError(f"自动红偏校正：文件夹中没有可读取的图片：{folder}") images: List[Image.Image] = [] for name in names: path = os.path.join(folder, name) with Image.open(path) as img: images.append(img.convert("RGB")) return images def _srgb_to_lab(image: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: rgb = image[..., :3].clamp(0.0, 1.0) linear = torch.where( rgb <= 0.04045, rgb / 12.92, ((rgb + 0.055) / 1.055).pow(2.4), ) red, green, blue = linear.unbind(dim=-1) x = (0.4124564 * red + 0.3575761 * green + 0.1804375 * blue) / D65_WHITE[0] y = 0.2126729 * red + 0.7151522 * green + 0.0721750 * blue z = (0.0193339 * red + 0.1191920 * green + 0.9503041 * blue) / D65_WHITE[2] def pivot(value: torch.Tensor) -> torch.Tensor: return torch.where( value > LAB_EPSILON, value.clamp_min(0.0).pow(1.0 / 3.0), (LAB_KAPPA * value + 16.0) / 116.0, ) fx, fy, fz = pivot(x), pivot(y), pivot(z) lightness = 116.0 * fy - 16.0 a = 500.0 * (fx - fy) b = 200.0 * (fy - fz) return lightness, a, b def _lab_to_srgb(lightness: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: fy = (lightness + 16.0) / 116.0 fx = fy + a / 500.0 fz = fy - b / 200.0 def inverse_pivot(value: torch.Tensor) -> torch.Tensor: cubed = value.pow(3.0) return torch.where(cubed > LAB_EPSILON, cubed, (116.0 * value - 16.0) / LAB_KAPPA) x = D65_WHITE[0] * inverse_pivot(fx) y = inverse_pivot(fy) z = D65_WHITE[2] * inverse_pivot(fz) red = 3.2404542 * x - 1.5371385 * y - 0.4985314 * z green = -0.9692660 * x + 1.8760108 * y + 0.0415560 * z blue = 0.0556434 * x - 0.2040259 * y + 1.0572252 * z linear = torch.stack((red, green, blue), dim=-1) positive = linear.clamp_min(0.0) srgb = torch.where( linear <= 0.0031308, 12.92 * linear, 1.055 * positive.pow(1.0 / 2.4) - 0.055, ) return srgb.clamp(0.0, 1.0) def _smoothstep(value: torch.Tensor) -> torch.Tensor: value = value.clamp(0.0, 1.0) return value * value * (3.0 - 2.0 * value) class O1keyAutoRedCast: """自动检测并移除商品图红偏，零 API 成本。支持单张和批量文件夹两种模式。""" @classmethod def INPUT_TYPES(cls): return { "required": { "强度": ( "FLOAT", { "default": 1.0, "min": 0.0, "max": 1.5, "step": 0.05, "tooltip": "1.0 为自动测得的完整校正量。", }, ), "最大校正量": ( "FLOAT", { "default": 8.0, "min": 0.0, "max": 20.0, "step": 0.5, "tooltip": "限制 Lab 红轴最大校正量，防止极端图片过度校色。", }, ), "高饱和保护": ( "FLOAT", { "default": 0.1, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "保护橙色、蓝色等高饱和区域；0 为统一白平衡，1 为最大保护。", }, ), "图片路径": ( "STRING", { "default": "", "multiline": False, "tooltip": ( "批量模式：填写文件夹路径后将处理其中所有图片，忽略上方图像输入。" "文件夹内图片必须尺寸一致。留空则使用图像输入端口。" ), }, ), }, "optional": { "图像": ( "IMAGE", { "tooltip": "单张模式输入；填写图片路径进入批量模式后可不连接。", }, ), "灰卡最低亮度": ( "FLOAT", { "default": 58.0, "min": 20.0, "max": 95.0, "step": 1.0, "tooltip": "灰卡候选区域的最低 Lab 亮度。", }, ), "灰卡最大色度": ( "FLOAT", { "default": 18.0, "min": 3.0, "max": 40.0, "step": 1.0, "tooltip": "灰卡候选区域允许的最大色度。", }, ), "seed": ( "INT", { "default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "step": 1, "control_after_generate": True, "tooltip": "ComfyUI 原生随机种子；改变 seed 可重新运行节点，校色结果由图像和校色参数决定。", }, ), }, } RETURN_TYPES = ("IMAGE", "MASK", "STRING") RETURN_NAMES = ("校正图像", "取样遮罩", "检测报告") FUNCTION = "correct" CATEGORY = "o1key/image" DESCRIPTION = ( "零成本自动检测并移除商品图红偏，不调用模型或网络 API。" "填写图片路径可批量处理整个文件夹；留空则处理连接的图像输入。" ) @torch.inference_mode() def correct( self, 图像: Optional[torch.Tensor] = None, 强度: float = 1.0, 最大校正量: float = 8.0, 高饱和保护: float = 0.1, 图片路径: str = "", seed: int = 0, 灰卡最低亮度: float = 58.0, 灰卡最大色度: float = 18.0, ): # --- 数据来源：文件夹 or 图像输入 --- if 图片路径 and 图片路径.strip(): pil_images = _load_folder_images(图片路径.strip()) # 校验所有图片尺寸一致（不同尺寸无法合并为批次张量） sizes = {img.size for img in pil_images} if len(sizes) > 1: raise ValueError( f"自动红偏校正：文件夹中的图片尺寸不统一 {sizes}，" "请确保所有图片宽高相同，或分批放入不同文件夹。" ) source = pil_to_tensor(pil_images) # (B, H, W, 3), float32, [0,1] print(f"[o1key 自动红偏校正] 批量模式：加载 {len(pil_images)} 张图片，尺寸 {pil_images[0].size}") else: if 图像 is None: raise ValueError( "自动红偏校正：请连接图像输入，或填写批量图片文件夹路径。" ) source = 图像 source_float = source.float() lightness, a, b = _srgb_to_lab(source_float) chroma = torch.hypot(a, b) corrected_a = a.clone() masks = [] report_lines = [] for index in range(source_float.shape[0]): sample_mask = (lightness[index] >= 灰卡最低亮度) & (chroma[index] <= 灰卡最大色度) minimum_pixels = max(1024, int(sample_mask.numel() * 0.001)) # 中性像素太少时放宽一档，避免深色背景图直接放弃校正 if int(sample_mask.sum().item()) < minimum_pixels: sample_mask = (lightness[index] >= max(40.0, 灰卡最低亮度 - 12.0)) & ( chroma[index] <= 灰卡最大色度 + 8.0 ) sample_count = int(sample_mask.sum().item()) masks.append(sample_mask.float()) if sample_count < minimum_pixels: report_lines.append(f"第 {index + 1} 张：中性取样不足，保持原图") continue measured_a = float(a[index][sample_mask].mean().item()) correction = max(0.0, min(float(最大校正量), measured_a + 0.2)) applied = correction * float(强度) saturation = _smoothstep((chroma[index] - 18.0) / 36.0) protection = 1.0 - float(高饱和保护) * saturation corrected_a[index] = a[index] - applied * protection sample_ratio = 100.0 * sample_count / sample_mask.numel() if applied > 0.01: report_lines.append( f"第 {index + 1} 张：检测红轴 {measured_a:+.2f}，" f"校正 {-applied:.2f}，取样 {sample_ratio:.1f}%" ) else: report_lines.append( f"第 {index + 1} 张：未检测到红偏，保持原图，取样 {sample_ratio:.1f}%" ) corrected_rgb = _lab_to_srgb(lightness, corrected_a, b) # 保留原始 alpha 通道（如有） if source_float.shape[-1] > 3: corrected = torch.cat((corrected_rgb, source_float[..., 3:]), dim=-1) else: corrected = corrected_rgb report = "\n".join(report_lines) print("[o1key 自动红偏校正] " + " | ".join(report_lines)) return (corrected.to(dtype=source.dtype), torch.stack(masks), report)