120 lines
4.8 KiB
Python
120 lines
4.8 KiB
Python
import torch
|
||
import numpy as np
|
||
from PIL import Image, ImageDraw
|
||
import folder_paths
|
||
import os
|
||
import uuid
|
||
|
||
|
||
def tensor2pil(image):
|
||
"""将 torch tensor 转换为 PIL Image"""
|
||
if len(image.shape) == 4:
|
||
image = image[0]
|
||
image_np = image.cpu().numpy()
|
||
if image_np.dtype != np.uint8:
|
||
if image_np.max() <= 1.0:
|
||
image_np = (image_np * 255).astype(np.uint8)
|
||
else:
|
||
image_np = np.clip(image_np, 0, 255).astype(np.uint8)
|
||
if len(image_np.shape) == 3 and image_np.shape[2] == 1:
|
||
image_np = image_np.squeeze(2)
|
||
return Image.fromarray(image_np)
|
||
|
||
|
||
def pil2tensor(image):
|
||
"""将 PIL Image 转换为 torch tensor"""
|
||
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
||
|
||
|
||
class ImageABCompare:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"ImageA": ("IMAGE",),
|
||
"ImageB": ("IMAGE",),
|
||
"split_direction": (["vertical", "horizontal"], {"default": "vertical"}), # 左右/上下
|
||
"split_ratio_when_output": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), # 分割比例
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("Image_output",)
|
||
FUNCTION = "compare_images"
|
||
CATEGORY = "🎨 图形比比看--ImageABCompareSee"
|
||
OUTPUT_NODE = True
|
||
|
||
def compare_images(self, ImageA, ImageB, split_direction="vertical", split_ratio_when_output=0.5):
|
||
print(f"[ImageABCompare] 生成导出用对比图 - 方向: {split_direction}, 比例: {split_ratio_when_output}")
|
||
|
||
# 转换图片（以image1尺寸为基准）
|
||
img1 = tensor2pil(ImageA)
|
||
img2 = tensor2pil(ImageB)
|
||
base_w, base_h = img1.size # 以图1尺寸为最终保存尺寸
|
||
print(f"[ImageABCompare] 基准尺寸（图1）: {base_w}x{base_h}，图2原始尺寸: {img2.size}")
|
||
|
||
# 保存原始图片到临时目录（供前端加载实时预览）
|
||
unique_id = str(uuid.uuid4())[:8]
|
||
img1_filename = f"ab_img1_{unique_id}.png"
|
||
img2_filename = f"ab_img2_{unique_id}.png"
|
||
output_dir = folder_paths.get_temp_directory()
|
||
img1.save(os.path.join(output_dir, img1_filename), format="PNG")
|
||
img2.save(os.path.join(output_dir, img2_filename), format="PNG")
|
||
|
||
# 按规则调整图2尺寸（用于导出）
|
||
if split_direction == "vertical":
|
||
img2_resized = img2.resize((base_w, base_h), Image.Resampling.LANCZOS)
|
||
else:
|
||
img2_resized = img2.resize((base_w, base_h), Image.Resampling.LANCZOS)
|
||
|
||
# 创建最终对比图画布（用于导出）
|
||
result_img = Image.new("RGB", (base_w, base_h))
|
||
draw = ImageDraw.Draw(result_img)
|
||
|
||
# 绘制对比图（用于导出）
|
||
if split_direction == "vertical":
|
||
split_x = int(base_w * split_ratio_when_output)
|
||
result_img.paste(img1.crop((0, 0, split_x, base_h)), (0, 0))
|
||
result_img.paste(img2_resized.crop((split_x, 0, base_w, base_h)), (split_x, 0))
|
||
draw.line([(split_x, 0), (split_x, base_h)], fill=(255, 0, 0), width=3)
|
||
else:
|
||
split_y = int(base_h * split_ratio_when_output)
|
||
result_img.paste(img1.crop((0, 0, base_w, split_y)), (0, 0))
|
||
result_img.paste(img2_resized.crop((0, split_y, base_w, base_h)), (0, split_y))
|
||
draw.line([(0, split_y), (base_w, split_y)], fill=(255, 0, 0), width=3)
|
||
|
||
# 生成导出用tensor
|
||
result_tensor = pil2tensor(result_img)
|
||
|
||
# 返回前端需要的所有参数（含实时比例）
|
||
return {
|
||
"ui": {
|
||
"images": [], # 不显示后端生成的静态图，避免覆盖前端画布
|
||
"img1_filename": [img1_filename], # 前端加载图1用
|
||
"img2_filename": [img2_filename], # 前端加载图2用
|
||
"split_direction": [split_direction],
|
||
"split_ratio_when_output": [split_ratio_when_output], # 实时返回滑块值
|
||
"base_size": [base_w, base_h]
|
||
},
|
||
"result": (result_tensor,) # 输出tensor可连接保存图像节点
|
||
}
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"ImageABCompare": ImageABCompare
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"ImageABCompare": "图形比比看--ImageABCompareSee"
|
||
}
|
||
|
||
if __name__ == "__main__":
|
||
print("=== 节点注册验证 ===")
|
||
print("注册的节点类:", NODE_CLASS_MAPPINGS.keys())
|
||
print("显示名称映射:", NODE_DISPLAY_NAME_MAPPINGS)
|
||
|
||
try:
|
||
node = ImageABCompare()
|
||
print("✅ 节点类实例化成功")
|
||
print("输入类型:", node.INPUT_TYPES())
|
||
except Exception as e:
|
||
print("❌ 节点类实例化失败:", str(e)) |