commit 705bfe12e9f646fad724cfc57e9bdb9e100f92b6 Author: kk8bit Date: Sun Oct 20 00:33:07 2024 +0800 v0.0.1 diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..e2f7218 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,12 @@ +# Changelog + +## [0.0.1] - 2024-10-19 +### Added +- 实现自定义保存图像节点。 +- 支持保存为 PNG 和 JPG 格式。 +- 添加 JPG 质量调整选项。 +- 支持嵌入作者和版权信息。 +- 支持选择 sRGB IEC61966-2.1 和 Adobe RGB (1998) 颜色配置文件。 +- 当选择 Adobe RGB 时,自动将图像转换为 Adobe RGB并保证色彩准确。 +- 支持保存元数据到 PNG 文件,包括 `prompt` 和 `extra_pnginfo`。 +- 自动生成唯一文件名。 \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..f159d5e --- /dev/null +++ b/README.md @@ -0,0 +1,24 @@ +# KayTool + +这是一个为 ComfyUI 开发的自定义节点实用工具包,在未来我会陆续为它增加功能 + +## 当前功能: + +- 选择保存图片的格式(PNG 或 JPG) +- 选择是否保存元数据(工作流信息) +- 支持两种颜色配置文件(sRGB IEC61966-2.1 和 Adobe RGB (1998)) +- 当选择 Adobe RGB 时,自动将图像转换为 Adobe RGB并保证色彩准确 +- 自定义 JPG 图片质量(0-100) +- 可以嵌入作者和版权信息 +- 自动生成唯一文件名 +- 在默认的`output`文件夹下,自动创建`custom_save_images`文件夹,以保存所有的图片 + +## 一些说明 + +- 要保存”JPG“格式,必须将"metadata“设置为"False" +- 当选择保存"metadata“(工作流信息)时,将自动切换为"PNG"格式,并忽略"JPG"格式的设置 +- 如果你需要更小容量但不损失肉眼可见的画质,可以讲"JPG"图片质量设置为 40,更低的数值将肉眼可见色块 + +## 安装与使用 + +- 将本项目克隆到你的 `ComfyUI/custom_nodes`目录下,并确保将 sRGB Profile.icc 和AdobeRGB1998.icc 文件放在 resources 目录中 \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ea8d026 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .custom_save_image import register_custom_node as register_save_node + +NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = register_save_node() \ No newline at end of file diff --git a/custom_save_image.py b/custom_save_image.py new file mode 100644 index 0000000..4c79dae --- /dev/null +++ b/custom_save_image.py @@ -0,0 +1,156 @@ +import os +import json +import numpy as np +from PIL import Image +from PIL.PngImagePlugin import PngInfo +from PIL import ExifTags, TiffImagePlugin, ImageCms + +class CustomSaveImage: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "filename_prefix": ("STRING", {"default": "Custom_Save_Image"}), + "save_metadata": ("BOOLEAN", {"default": True}), + "color_profile": (["sRGB IEC61966-2.1", "Adobe RGB (1998)"], {"default": "sRGB IEC61966-2.1"}), # 将 color_profile 移到 save_metadata 下面 + "format": (["PNG", "JPG"], {"default": "PNG"}), + "jpg_quality": ("INT", {"default": 95, "min": 0, "max": 100}), + "author": ("STRING", {"default": ""}), + "copyright_info": ("STRING", {"default": ""}) + }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + + RETURN_TYPES = () + FUNCTION = "save_images" + + OUTPUT_NODE = True + + CATEGORY = "KayTool" + + def save_images(self, images, filename_prefix="Custom_Save_Image", save_metadata=True, format="PNG", jpg_quality=95, author="", copyright_info="", color_profile="sRGB IEC61966-2.1", prompt=None, extra_pnginfo=None): + output_dir = self.get_output_directory() + os.makedirs(output_dir, exist_ok=True) # 确保目录存在 + + results = [] # 存储图片保存路径以供UI预览使用 + + # 动态获取 sRGB 和 AdobeRGB 的 ICC profile 路径 + srgb_profile_path = os.path.join(os.path.dirname(__file__), 'resources', 'sRGB Profile.icc') + adobergb_profile_path = os.path.join(os.path.dirname(__file__), 'resources', 'AdobeRGB1998.icc') + + for image in images: + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + + # 处理颜色配置文件 + if color_profile == "Adobe RGB (1998)": + img = self.convert_to_adobe_rgb(img, srgb_profile_path, adobergb_profile_path) + icc_profile_path = adobergb_profile_path + else: + icc_profile_path = srgb_profile_path + + # 加载选择的 ICC profile + icc_profile = self.load_icc_profile(icc_profile_path) + + # 如果 save_metadata 为 True,默认保存为 PNG + if save_metadata: + metadata = PngInfo() + + # 无论 save_metadata,始终嵌入作者和版权信息 + if author: + metadata.add_text("Author", author) + if copyright_info: + metadata.add_text("Copyright", copyright_info) + + # 嵌入额外元数据 + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for k, v in extra_pnginfo.items(): + metadata.add_text(k, json.dumps(v)) + + # 保存为 PNG 格式,嵌入所选颜色配置文件 + filename = f"{filename_prefix}_{self.get_unique_filename()}.png" + full_output_path = os.path.join(output_dir, filename) + img.save(full_output_path, pnginfo=metadata, icc_profile=icc_profile) + + # 如果 save_metadata 为 False,保存为指定的格式 + else: + if format == "PNG": + metadata = PngInfo() + + # 嵌入作者和版权信息 + if author: + metadata.add_text("Author", author) + if copyright_info: + metadata.add_text("Copyright", copyright_info) + + # 保存为 PNG 格式 + filename = f"{filename_prefix}_{self.get_unique_filename()}.png" + full_output_path = os.path.join(output_dir, filename) + img.save(full_output_path, pnginfo=metadata, icc_profile=icc_profile) + + elif format == "JPG": + exif_data = img.getexif() + + # 嵌入作者和版权信息 + if author: + exif_data[0x013B] = author # EXIF中的作者字段(Artist) + if copyright_info: + exif_data[0x8298] = copyright_info # EXIF中的版权字段(Copyright) + + exif_bytes = exif_data.tobytes() + + # 保存为 JPG 格式,附加 EXIF 和 ICC profile + filename = f"{filename_prefix}_{self.get_unique_filename()}.jpg" + full_output_path = os.path.join(output_dir, filename) + img.save(full_output_path, quality=jpg_quality, exif=exif_bytes, icc_profile=icc_profile) + + # 添加保存的图片信息用于UI预览 + results.append({ + "filename": filename, + "subfolder": "Custom_Save_Image", + "type": "output", + }) + + # 返回给UI的图片预览信息 + return {"ui": {"images": results}, "status": "Images saved successfully"} + + def convert_to_adobe_rgb(self, img, srgb_profile_path, adobergb_profile_path): + """ + 将图像从 sRGB 转换为 Adobe RGB + """ + srgb_profile = ImageCms.getOpenProfile(srgb_profile_path) + adobergb_profile = ImageCms.getOpenProfile(adobergb_profile_path) + img = ImageCms.profileToProfile(img, srgb_profile, adobergb_profile) + return img + + def load_icc_profile(self, path): + """ + 加载 ICC profile,如果加载失败,抛出异常 + """ + try: + with open(path, "rb") as f: + return f.read() + except FileNotFoundError: + raise Exception(f"ICC profile not found at: {path}") + + def get_unique_filename(self): + """ + 生成基于时间戳的唯一文件名 + """ + import time + return f"{int(time.time() * 1000)}" + + def get_output_directory(self): + """ + 返回输出目录 + """ + return os.path.join(os.getcwd(), "output", "Custom_Save_Image") + +# 这个函数用于注册节点 +def register_custom_node(): + NODE_CLASS_MAPPINGS = {"CustomSaveImage": CustomSaveImage} + NODE_DISPLAY_NAME_MAPPINGS = {"CustomSaveImage": "Custom Save Image"} + return NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS \ No newline at end of file diff --git a/resources/AdobeRGB1998.icc b/resources/AdobeRGB1998.icc new file mode 100644 index 0000000..a79f576 Binary files /dev/null and b/resources/AdobeRGB1998.icc differ diff --git a/resources/sRGB Profile.icc b/resources/sRGB Profile.icc new file mode 100644 index 0000000..7f9d18d Binary files /dev/null and b/resources/sRGB Profile.icc differ