Jony
ba920f2b66
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
290 lines
11 KiB
Python
290 lines
11 KiB
Python
"""
|
||
自动红偏校正。
|
||
|
||
纯 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)
|