142 lines
5.5 KiB
Python
142 lines
5.5 KiB
Python
import numpy as np
|
|
import torch
|
|
from skimage.color import rgb2lab
|
|
from skimage.util import img_as_ubyte
|
|
from PIL import Image
|
|
import matplotlib.pyplot as plt
|
|
|
|
class lightdarkjudgment:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image1": ("IMAGE", ),
|
|
"image2": ("IMAGE", ),
|
|
"threshold": ("FLOAT", {"default": 32.0, "min": 0.1, "max": 100.0, "step": 0.01}),
|
|
"brightness1": ("FLOAT", {"default": 1, "min": 0.0, "max": 3, "step": 0.01}),
|
|
"contrast1": ("FLOAT", {"default": 1, "min": 0.0, "max": 3, "step": 0.01}),
|
|
"saturation1": ("FLOAT", {"default": 1, "min": 0.0, "max": 3, "step": 0.01}),
|
|
"brightness2": ("FLOAT", {"default": 1, "min": 0.0, "max": 3, "step": 0.01}),
|
|
"contrast2": ("FLOAT", {"default": 1, "min": 0.0, "max": 3, "step": 0.01}),
|
|
"saturation2": ("FLOAT", {"default": 1, "min": 0.0, "max": 3, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "FLOAT", "FLOAT", "FLOAT", "FLOAT")
|
|
RETURN_NAMES = ("IMAGE", "brightness", "contrast", "saturation", "color diff",)
|
|
FUNCTION = "lightordarkimage"
|
|
CATEGORY = "MoneyMaker😺"
|
|
|
|
|
|
|
|
|
|
def calculate_color_difference_in_diff_area(self, img1, img2):
|
|
"""
|
|
计算两张图像的色差，忽略透明区域
|
|
|
|
参数:
|
|
img1 (numpy array): 第一张图像
|
|
img2 (numpy array): 第二张图像
|
|
|
|
返回:
|
|
float: 差异区域的平均色差
|
|
"""
|
|
# 提取透明度通道
|
|
alpha1 = img1[:, :, 3] if img1.shape[2] == 4 else np.ones(img1.shape[:2])
|
|
alpha2 = img2[:, :, 3] if img2.shape[2] == 4 else np.ones(img2.shape[:2])
|
|
|
|
# 创建遮罩，只保留不透明区域
|
|
mask = (alpha1 > 0.1) & (alpha2 > 0.1)
|
|
|
|
# 转换图像为Lab颜色空间
|
|
img1_lab = rgb2lab(img1[:, :, :3])
|
|
img2_lab = rgb2lab(img2[:, :, :3])
|
|
|
|
# 计算每个像素的绝对差异
|
|
diff_lab = np.abs(img1_lab - img2_lab)
|
|
|
|
# 计算整体差异的掩码 (即差异区域的掩码)
|
|
diff_mask = np.any(diff_lab > 1, axis=-1) & mask # 选择阈值为1，阈值可以根据需要调整
|
|
|
|
# 在掩码区域内计算色差
|
|
diff_in_diff_area_lab = np.sqrt(np.sum(diff_lab ** 2, axis=-1))
|
|
diff_in_diff_area_lab[~diff_mask] = 0 # 仅保留差异区域的色差
|
|
|
|
# 如果是一个极小值，则设置为0值
|
|
if np.isnan(np.mean(diff_in_diff_area_lab[diff_mask])):
|
|
return 0.0
|
|
else:
|
|
return np.mean(diff_in_diff_area_lab[diff_mask])
|
|
|
|
def lightordarkimage(self, image1, image2, threshold, brightness1, contrast1, saturation1, brightness2, contrast2, saturation2):
|
|
"""
|
|
根据色差阈值判断并返回对应的图像及其属性
|
|
|
|
参数:
|
|
image1 (torch.Tensor): 第一张图像
|
|
image2 (torch.Tensor): 第二张图像
|
|
threshold (float): 色差阈值
|
|
brightness1 (float): 第一张图像的亮度
|
|
contrast1 (float): 第一张图像的对比度
|
|
saturation1 (float): 第一张图像的饱和度
|
|
brightness2 (float): 第二张图像的亮度
|
|
contrast2 (float): 第二张图像的对比度
|
|
saturation2 (float): 第二张图像的饱和度
|
|
|
|
返回:
|
|
tuple: 根据色差阈值返回的图像及其属性
|
|
"""
|
|
B, H, W, C = image1.shape
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
image1_tensor = image1.permute(0, 3, 1, 2).clone().to(device)
|
|
image2_tensor = image2.permute(0, 3, 1, 2).clone().to(device)
|
|
|
|
if image1.shape[1:] != image2.shape[1:]:
|
|
image2_tensor = torch.nn.functional.interpolate(image2_tensor, size=(H, W), mode="bilinear")
|
|
|
|
if image2.shape[0] < B:
|
|
image2_tensor = image2_tensor[0].unsqueeze(0).repeat(B, 1, 1, 1)
|
|
|
|
# 转换为numpy格式并计算色差
|
|
image1_np = image1_tensor.permute(0, 2, 3, 1).cpu().numpy()[0]
|
|
image2_np = image2_tensor.permute(0, 2, 3, 1).cpu().numpy()[0]
|
|
|
|
# 计算色差
|
|
color_diff = self.calculate_color_difference_in_diff_area(image1_np, image2_np)
|
|
print(f"Color Difference in Difference Area: {color_diff}")
|
|
|
|
|
|
# 根据色差阈值判断并返回对应的图像和属性
|
|
if color_diff < threshold:
|
|
tensor_out = image1_tensor
|
|
brightness = brightness1
|
|
contrast = contrast1
|
|
saturation = saturation1
|
|
else:
|
|
tensor_out = image2_tensor
|
|
brightness = brightness2
|
|
contrast = contrast2
|
|
saturation = saturation2
|
|
|
|
tensor_out = torch.clamp(tensor_out, 0, 1)
|
|
tensor_out = tensor_out.permute(0, 2, 3, 1).cpu().float()
|
|
|
|
return tensor_out, brightness, contrast, saturation, color_diff
|
|
|
|
# A dictionary that contains all nodes you want to export with their names
|
|
# NOTE: names should be globally unique,一行不需要逗号，多行需要逗号
|
|
# 冒号后为函数名
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Light or Dark": lightdarkjudgment
|
|
}
|
|
|
|
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Light or Dark": "Light or Dark"
|
|
# After the colon is the name displayed on the node. Multiple lines correspond to multiple nodes
|
|
}
|