89 lines
3.0 KiB
Python
89 lines
3.0 KiB
Python
import os
|
|
import numpy as np
|
|
from PIL import Image, ImageOps
|
|
from torch import Tensor
|
|
from typing import Tuple
|
|
from .utils.image_to_tensor import image_to_tensor
|
|
|
|
class LoadLastImage:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"initial_image": ("IMAGE",),
|
|
"save_initial_image": ("BOOLEAN", {"default": True}),
|
|
"directory_path": ("STRING", {"multiline": False}),
|
|
"sort": (["created", "modified"],)
|
|
}
|
|
}
|
|
|
|
CATEGORY = "TrollSuite/image"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "LoadLastImage"
|
|
|
|
def LoadLastImage(self, initial_image: Tensor, save_initial_image: bool, directory_path: str, sort: str) -> Tuple[Tensor]:
|
|
image = initial_image
|
|
if os.path.isabs(directory_path):
|
|
directory_path = os.path.abspath(directory_path)
|
|
else:
|
|
directory_path = os.path.join("./output", os.path.relpath(directory_path))
|
|
|
|
# create directory
|
|
self.mkdirs(directory_path)
|
|
|
|
# store initial image if dir is empty
|
|
if save_initial_image and self.is_directory_empty(directory_path):
|
|
for index, image in enumerate(image):
|
|
i = 255. * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
img.save(os.path.join(directory_path, f"{index:05}_.png"), pnginfo=None, compress_level=0)
|
|
|
|
filepath = self.get_last_filepath(directory_path, sort)
|
|
if filepath:
|
|
image = self.convert_to_tensor(filepath)
|
|
return (image,)
|
|
|
|
def get_last_filepath(self, directory_path:str, sort: str) -> str:
|
|
"""
|
|
Get the latest image file in the given directory
|
|
"""
|
|
# Resolve relative directories and remove trailing slashes
|
|
directory_path = os.path.abspath(directory_path)
|
|
|
|
# get filepaths of all files and dirs in the given dir
|
|
valid_files = [os.path.join(directory_path, filename) for filename in os.listdir(directory_path)]
|
|
|
|
if not valid_files:
|
|
return None
|
|
|
|
if sort == "modified":
|
|
key = os.path.getmtime
|
|
else:
|
|
key = os.path.getctime
|
|
|
|
return max(valid_files, key=key)
|
|
|
|
def mkdirs(self, directory_path: str):
|
|
"""
|
|
Creates the specified directory if it does not exist.
|
|
"""
|
|
if not os.path.exists(directory_path):
|
|
try:
|
|
os.makedirs(directory_path)
|
|
except OSError:
|
|
print ("Creation of the directory %s failed" % directory_path)
|
|
|
|
def is_directory_empty(self, directory_path:str):
|
|
return not bool(os.listdir(directory_path))
|
|
|
|
def convert_to_tensor(self, image_path: str) -> Tensor:
|
|
i = Image.open(image_path)
|
|
i = ImageOps.exif_transpose(i)
|
|
image = i.convert("RGB")
|
|
return image_to_tensor(image)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(self):
|
|
return float("NaN")
|