diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..96fc161 --- /dev/null +++ b/__init__.py @@ -0,0 +1,13 @@ +from .canvas_node import CanvasView + +NODE_CLASS_MAPPINGS = { + "CanvasView": CanvasView +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "CanvasView": "Canvas View" +} + +WEB_DIRECTORY = "./js" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] \ No newline at end of file diff --git a/canvas_node.py b/canvas_node.py new file mode 100644 index 0000000..45579a9 --- /dev/null +++ b/canvas_node.py @@ -0,0 +1,61 @@ +from PIL import Image, ImageOps +import hashlib +import torch +import numpy as np +import folder_paths +from server import PromptServer + +class CanvasView: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "canvas_image": ("STRING", {"default": "canvas_image.png"}), + "trigger": ("INT", {"default": 0, "min": 0, "max": 99999999, "step": 1}) + }, + "hidden": { + "unique_id": "UNIQUE_ID" + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "mask") + FUNCTION = "process_canvas_image" + CATEGORY = "ycnode" + + def process_canvas_image(self, canvas_image, trigger, unique_id): + try: + # 读取保存的画布图像和遮罩 + path_image = folder_paths.get_annotated_filepath(canvas_image) + path_mask = folder_paths.get_annotated_filepath(canvas_image.replace('.png', '_mask.png')) + + # 处理主图像 + i = Image.open(path_image) + i = ImageOps.exif_transpose(i) + if i.mode not in ['RGB', 'RGBA']: + i = i.convert('RGB') + image = np.array(i).astype(np.float32) / 255.0 + if i.mode == 'RGBA': + rgb = image[..., :3] + alpha = image[..., 3:] + image = rgb * alpha + (1 - alpha) * 0.5 + + # 处理遮罩图像 + try: + mask = Image.open(path_mask).convert('L') + mask = np.array(mask).astype(np.float32) / 255.0 + mask = torch.from_numpy(mask)[None,] + except: + # 如果没有遮罩文件,创建空白遮罩 + mask = torch.zeros((1, image.shape[0], image.shape[1]), dtype=torch.float32) + + # 转换为tensor + image = torch.from_numpy(image)[None,] + + return (image, mask) + except Exception as e: + print(f"Error processing canvas image: {str(e)}") + # 返回白色图像和空白遮罩 + blank = np.ones((512, 512, 3), dtype=np.float32) + blank_mask = np.zeros((512, 512), dtype=np.float32) + return (torch.from_numpy(blank)[None,], torch.from_numpy(blank_mask)[None,]) \ No newline at end of file