add auxiliary tool

This commit is contained in:
yueliang
2024-10-09 18:25:56 +08:00
parent 67536b0b1d
commit ff4153667d
2 changed files with 61 additions and 0 deletions
+58
View File
@@ -0,0 +1,58 @@
import os
import torch
from PIL import ImageOps
import comfy
from PIL import Image
import numpy as np
class AnyType(str):
def __ne__(self, __value: object) -> bool:
return False
any_type = AnyType("*")
class ImagePath2Tensor:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image_paths": (any_type,),
},
}
RETURN_TYPES = ("IMAGE","INT")
RETURN_NAMES = ("image","image_count")
FUNCTION = "image_path_to_tensor"
OUTPUT_NODE = True
CATEGORY = "🔥FFmpeg/auxiliary tool"
def image_path_to_tensor(self, image_paths):
try:
#['D:\\Cache\\222\\frame_00000121.png', 'D:\\Cache\\222\\frame_00000122.png']
images = []
for image_path in image_paths:
# Open and process the image
with Image.open(image_path) as img:
img = ImageOps.exif_transpose(img).convert("RGB")
# Directly convert to torch tensor without numpy intermediate
image_tensor = torch.from_numpy(np.array(np.array(img).astype(np.float32) / 255.0))[None,]
images.append(image_tensor)
if len(images) == 0:
raise ValueError("No images loaded successfully.")
if len(images) == 1:
return (images[0], 1)
elif len(images) > 1:
image1 = images[0]
for image2 in images[1:]:
if image1.shape[1:] != image2.shape[1:]:
image2 = comfy.utils.common_upscale(image2.movedim(-1, 1), image1.shape[2], image1.shape[1], "bilinear", "center").movedim(1, -1)
image1 = torch.cat((image1, image2), dim=0)
return (image1, len(images))
except Exception as e:
raise ValueError(e)
+3
View File
@@ -6,6 +6,7 @@ from .nodes.videoFlip import *
from .nodes.extractAudio import *
from .nodes.loadImageFromDir import *
from .nodes.imageCopy import *
from .nodes.imagePath2Tensor import *
NODE_CLASS_MAPPINGS = {
"Video2Frames": Video2Frames,
@@ -16,6 +17,7 @@ NODE_CLASS_MAPPINGS = {
"ExtractAudio": ExtractAudio,
"LoadImageFromDir": LoadImageFromDir,
"ImageCopy": ImageCopy,
"ImagePath2Tensor": ImagePath2Tensor,
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -28,4 +30,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"ExtractAudio": "🔥ExtractAudio",
"LoadImageFromDir": "🔥LoadImageFromDir",
"ImageCopy": "🔥ImageCopy",
"ImagePath2Tensor": "🔥ImagePath2Tensor",
}