59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
import os
|
||
import uuid
|
||
from PIL import Image
|
||
import numpy as np
|
||
import torch
|
||
|
||
class SaveImagePath:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"image_path": ("IMAGE", {"forceInput": True}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
FUNCTION = "load"
|
||
CATEGORY = "不忘科技-自定义节点🚩"
|
||
|
||
def load(self, image_path):
|
||
# 如果是torch.Tensor类型，转换为numpy数组
|
||
if isinstance(image_path, torch.Tensor):
|
||
image_path = image_path.cpu().numpy()
|
||
|
||
# 确保数据类型为uint8且在0 - 255范围内
|
||
if image_path.dtype != np.uint8:
|
||
image_path = np.clip(image_path, 0, 255).astype(np.uint8)
|
||
|
||
# 去除多余的维度，如果形状是(1, 1, height, width, channels)或(1, height, width, channels)等情况
|
||
while len(image_path.shape) > 3:
|
||
image_path = image_path.squeeze(0)
|
||
|
||
# 如果是单通道图像，转换为3通道
|
||
if len(image_path.shape) == 2:
|
||
image_path = np.stack([image_path] * 3, axis=-1)
|
||
|
||
# 如果是通道优先格式 (C, H, W)，转换为通道最后格式 (H, W, C)
|
||
elif len(image_path.shape) == 3 and image_path.shape[0] <= 4:
|
||
image_path = np.transpose(image_path, (1, 2, 0))
|
||
|
||
pil_image = Image.fromarray(image_path)
|
||
|
||
output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
|
||
if not os.path.exists(output_dir):
|
||
os.makedirs(output_dir)
|
||
|
||
file_name = "%s.jpg" % str(uuid.uuid4())
|
||
p = os.path.join(output_dir, file_name)
|
||
pil_image.save(p)
|
||
|
||
return (p,)
|
||
# 节点类定义结束，以下是用于注册节点的字典结构（通常在实际使用中由ComfyUI等框架来解析和注册）
|
||
NODE_CLASS_MAPPINGS = {
|
||
"SaveImagePath": SaveImagePath
|
||
}
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"SaveImagePath": "保存图片路径"
|
||
}
|