74 lines
2.8 KiB
Python
74 lines
2.8 KiB
Python
import os
|
|
from PIL import Image, ImageOps, ImageSequence
|
|
import numpy as np
|
|
import torch
|
|
# 导入 folder_paths 模块
|
|
import folder_paths # type: ignore
|
|
from .hua_icons import icons
|
|
|
|
class Go_to_image:
|
|
_color_channels = 3 # 假设RGB颜色通道数为3
|
|
@classmethod
|
|
def INPUT_TYPES(s): #定义输入类型
|
|
input_dir = folder_paths.get_input_directory() # 获取输入目录
|
|
files = sorted(os.listdir(input_dir)) # 获取输入目录下的文件列表，并按字母排序
|
|
return {
|
|
"required": { # 必需输入
|
|
"image": (files, {"image_upload": True}), # 图像文件名，允许上传
|
|
"pos_text": ("STRING", {"multiline": True, "default": "positive text"}), # 正面文本，默认为"positive text"
|
|
|
|
"images": ("IMAGE", ),
|
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "CONDITIONING") #返回类型，为3个输出口
|
|
|
|
|
|
FUNCTION = "load_image" #函数名称
|
|
|
|
CATEGORY = icons.get("hua_boy_one") #传递一级类显示在树列表
|
|
|
|
def load_image(self, image):
|
|
image_path = folder_paths.get_annotated_filepath(image)
|
|
img = Image.open(image_path)
|
|
output_images = []
|
|
output_masks = []
|
|
for i in ImageSequence.Iterator(img):
|
|
i = ImageOps.exif_transpose(i)
|
|
if i.mode == 'I':
|
|
i = i.point(lambda i: i * (1 / 255))
|
|
image = i.convert("RGB")
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
if 'A' in i.getbands():
|
|
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
|
|
mask = 1. - torch.from_numpy(mask)
|
|
else:
|
|
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
|
|
output_images.append(image)
|
|
output_masks.append(mask.unsqueeze(0))
|
|
|
|
if len(output_images) > 1:
|
|
output_image = torch.cat(output_images, dim=0)
|
|
output_mask = torch.cat(output_masks, dim=0)
|
|
else:
|
|
output_image = output_images[0]
|
|
output_mask = output_masks[0]
|
|
|
|
return (output_image, output_mask)
|
|
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"brucelee": Go_to_image # 赋值类的名称
|
|
}
|
|
|
|
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
|
NODE_DISPLAY_NAME_MAPPINGS = { #这是个字典，它将节点名称映射到人类容易读取的名字，可以通过 NODE_DISPLAY_NAME_MAPPINGS["Hua"] 访问到 "Demo Node" （第二级子目录）
|
|
"brucelee": "思维导图" # 相当于传递上面的："brucelee"= Go_to_image = "brucelee": "HUA_go_image"
|
|
}
|