84 lines
2.3 KiB
Python
84 lines
2.3 KiB
Python
import glob
|
|
import os
|
|
|
|
import torch
|
|
from PIL import Image
|
|
import torchvision
|
|
from torchvision.transforms import functional as TF
|
|
|
|
from . import utils
|
|
|
|
|
|
class RandomImageLoader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"directory": ("STRING", {"multiline": False, "default": "./input"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "doit"
|
|
OUTPUT_NODE = False
|
|
CATEGORY = "utils"
|
|
|
|
def doit(self, directory, seed):
|
|
directory = utils.get_directory(directory)
|
|
print("RandomImageLoader: directory = %s" % directory)
|
|
|
|
files = (glob.glob(os.path.join(directory, "*.png")) +
|
|
glob.glob(os.path.join(directory, "*.jpg")) +
|
|
glob.glob(os.path.join(directory, "*.jpeg")))
|
|
|
|
if len(files) == 0:
|
|
raise ValueError("Specified directory does not contain any image files")
|
|
|
|
file = files[seed % len(files)]
|
|
print("RandomImageLoader: load %s; in %d files" % (file, len(files)))
|
|
|
|
im0 = Image.open(file)
|
|
im1 = TF.to_tensor(im0.convert("RGBA"))
|
|
im1[:3, im1[3, :, :] == 0] = 0
|
|
|
|
images = torch.stack([im1])
|
|
images = images.permute(0, 2, 3, 1)
|
|
images = images[:, :, :, :3]
|
|
return (images,)
|
|
|
|
|
|
class SaveImageToDirectory:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"directory": ("STRING", {"multiline": False, "default": "./output"}),
|
|
"image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "doit"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "utils"
|
|
|
|
def doit(self, directory, image):
|
|
directory = utils.get_directory(directory)
|
|
print("SaveImageToDirectory: directory = %s" % directory)
|
|
|
|
next_index = utils.find_next_index(directory)
|
|
file_name = os.path.join(directory, "out.%06d.png" % next_index)
|
|
print("SaveImageToDirectory: save to %s" % file_name)
|
|
|
|
im0 = torchvision.transforms.functional.to_pil_image(torch.permute(image[0], (2, 0, 1)))
|
|
im0.save(file_name)
|
|
|
|
return (image,)
|