Added utility nodes
This commit is contained in:
Fillip
@@ -69,6 +69,8 @@ from .nodes.FL_BulkPDFLoader import FL_BulkPDFLoader
|
||||
from .nodes.FL_SaveAndDisplayImage import FL_SaveAndDisplayImage
|
||||
from .nodes.FL_OllamaCaptioner import FL_OllamaCaptioner
|
||||
from .nodes.FL_ImageAdjuster import FL_ImageAdjuster
|
||||
from .nodes.FL_Caption_Saver_V2 import FL_CaptionSaver_V2
|
||||
from .nodes.FL_PathTypeChecker import FL_PathTypeChecker
|
||||
|
||||
|
||||
|
||||
@@ -145,6 +147,8 @@ NODE_CLASS_MAPPINGS = {
|
||||
"FL_SaveAndDisplayImage": FL_SaveAndDisplayImage,
|
||||
"FL_OllamaCaptioner": FL_OllamaCaptioner,
|
||||
"FL_ImageAdjuster": FL_ImageAdjuster,
|
||||
"FL_CaptionSaver_V2": FL_CaptionSaver_V2,
|
||||
"FL_PathTypeChecker": FL_PathTypeChecker,
|
||||
|
||||
}
|
||||
|
||||
@@ -220,6 +224,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FL_SaveAndDisplayImage": "FL Save And Display Image",
|
||||
"FL_OllamaCaptioner": "FL Ollama Captioner by Cosmic",
|
||||
"FL_ImageAdjuster": "FL_ImageAdjuster",
|
||||
"FL_CaptionSaver_V2": "FL Caption Saver V2",
|
||||
"FL_PathTypeChecker": "FL Path Type Checker",
|
||||
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
import os
|
||||
import re
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
from comfy.utils import ProgressBar
|
||||
|
||||
|
||||
class FL_CaptionSaver_V2:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"input_type": (["Image Input", "Directory Input"], {"default": "Image Input"}),
|
||||
"caption_input_type": (["Single Caption", "Multiple Captions"], {"default": "Single Caption"}),
|
||||
"folder_name": ("STRING", {"default": "output_folder"}),
|
||||
"overwrite": ("BOOLEAN", {"default": True}),
|
||||
"downsize_factor": ([1, 2, 3], {"default": 1})
|
||||
},
|
||||
"optional": {
|
||||
"images": ("IMAGE", {}),
|
||||
"input_directory": ("STRING", {"default": ""}),
|
||||
"single_caption": ("STRING", {"default": "Your caption here"}),
|
||||
"multiple_captions": ("STRING", {"multiline": True, "default": ""})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
FUNCTION = "save_images_with_captions"
|
||||
CATEGORY = "🏵️Fill Nodes/Captioning"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
def sanitize_text(self, text):
|
||||
return re.sub(r'[^a-zA-Z0-9\s.,!?-]', '', text)
|
||||
|
||||
def save_images_with_captions(self, input_type, caption_input_type, folder_name, overwrite, downsize_factor,
|
||||
images=None, input_directory=None, single_caption="", multiple_captions=""):
|
||||
os.makedirs(folder_name, exist_ok=True)
|
||||
|
||||
if input_type == "Image Input" and images is not None:
|
||||
image_list = images
|
||||
use_original_names = False
|
||||
elif input_type == "Directory Input" and input_directory:
|
||||
image_list = [f for f in os.listdir(input_directory) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))]
|
||||
use_original_names = True
|
||||
else:
|
||||
return ("No valid input provided.",)
|
||||
|
||||
if caption_input_type == "Single Caption":
|
||||
captions = [self.sanitize_text(single_caption)] * len(image_list)
|
||||
else:
|
||||
captions = [self.sanitize_text(cap.strip()) for cap in multiple_captions.split('\n') if cap.strip()]
|
||||
if len(captions) < len(image_list):
|
||||
captions.extend([captions[-1]] * (len(image_list) - len(captions)))
|
||||
elif len(captions) > len(image_list):
|
||||
captions = captions[:len(image_list)]
|
||||
|
||||
saved_files = []
|
||||
pbar = ProgressBar(len(image_list))
|
||||
for i, (image_item, caption) in enumerate(zip(image_list, captions)):
|
||||
if use_original_names:
|
||||
base_name = os.path.splitext(image_item)[0]
|
||||
image_path = os.path.join(input_directory, image_item)
|
||||
image = Image.open(image_path)
|
||||
else:
|
||||
base_name = f"image_{i}"
|
||||
image_np = image_item.cpu().numpy()
|
||||
image_np = self.process_image_tensor(image_np)
|
||||
image = Image.fromarray(image_np)
|
||||
|
||||
# Downsize the image
|
||||
if downsize_factor > 1:
|
||||
new_size = (image.width // downsize_factor, image.height // downsize_factor)
|
||||
image = image.resize(new_size, Image.LANCZOS)
|
||||
|
||||
image_file_name = f"{folder_name}/{base_name}.png"
|
||||
text_file_name = f"{folder_name}/{base_name}.txt"
|
||||
|
||||
if not overwrite:
|
||||
image_file_name, text_file_name = self.get_unique_filenames(folder_name, base_name)
|
||||
|
||||
image.save(image_file_name)
|
||||
saved_files.append(image_file_name)
|
||||
|
||||
with open(text_file_name, "w") as text_file:
|
||||
text_file.write(caption)
|
||||
|
||||
pbar.update_absolute(i)
|
||||
|
||||
return (f"Saved {len(saved_files)} images (downsized by factor {downsize_factor}) and captions in '{folder_name}'",)
|
||||
|
||||
def process_image_tensor(self, image_np):
|
||||
if image_np.shape[0] == 1:
|
||||
image_np = np.squeeze(image_np, axis=0)
|
||||
if len(image_np.shape) == 2:
|
||||
image_np = np.stack((image_np,) * 3, axis=-1)
|
||||
elif image_np.shape[2] == 1:
|
||||
image_np = np.repeat(image_np, 3, axis=2)
|
||||
return (image_np * 255).clip(0, 255).astype(np.uint8)
|
||||
|
||||
def get_unique_filenames(self, folder_name, base_name):
|
||||
counter = 1
|
||||
image_file_name = f"{folder_name}/{base_name}.png"
|
||||
text_file_name = f"{folder_name}/{base_name}.txt"
|
||||
while os.path.exists(image_file_name) or os.path.exists(text_file_name):
|
||||
image_file_name = f"{folder_name}/{base_name}_{counter}.png"
|
||||
text_file_name = f"{folder_name}/{base_name}_{counter}.txt"
|
||||
counter += 1
|
||||
return image_file_name, text_file_name
|
||||
@@ -45,17 +45,17 @@ class FL_ImageCaptionSaver:
|
||||
|
||||
# Convert tensor to numpy array
|
||||
image_np = image_tensor.cpu().numpy()
|
||||
|
||||
|
||||
# Ensure the image is in the correct shape (height, width, channels)
|
||||
if image_np.shape[0] == 1: # If the first dimension is 1, squeeze it
|
||||
image_np = np.squeeze(image_np, axis=0)
|
||||
|
||||
|
||||
# If the image is grayscale (2D), convert to RGB
|
||||
if len(image_np.shape) == 2:
|
||||
image_np = np.stack((image_np,) * 3, axis=-1)
|
||||
elif image_np.shape[2] == 1: # If it's (height, width, 1)
|
||||
image_np = np.repeat(image_np, 3, axis=2)
|
||||
|
||||
|
||||
# Ensure values are in 0-255 range
|
||||
image_np = (image_np * 255).clip(0, 255).astype(np.uint8)
|
||||
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
import os
|
||||
import pathlib
|
||||
|
||||
class FL_PathTypeChecker:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"input_path": ("STRING", {"default": "", "multiline": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("PATH",)
|
||||
FUNCTION = "check_path_type"
|
||||
CATEGORY = "🏵️Fill Nodes/Utils"
|
||||
|
||||
def check_path_type(self, input_path):
|
||||
input_path = input_path.strip() # Remove leading/trailing whitespace
|
||||
|
||||
if not input_path:
|
||||
return ("Empty path provided.",)
|
||||
|
||||
path = pathlib.Path(input_path)
|
||||
|
||||
if path.is_absolute():
|
||||
return ("Absolute path",)
|
||||
elif path.is_relative_to(pathlib.Path.cwd()):
|
||||
return ("Relative path",)
|
||||
else:
|
||||
# Check if it might be a valid relative path
|
||||
try:
|
||||
path.relative_to(".")
|
||||
return ("Relative path",)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# If it's not recognized as absolute or relative, it might be invalid or a special case
|
||||
if os.path.splitdrive(input_path)[0]:
|
||||
return ("Drive-specific path",)
|
||||
elif input_path.startswith('//') or input_path.startswith('\\\\'):
|
||||
return ("UNC path",)
|
||||
elif '://' in input_path:
|
||||
return ("URL-like path",)
|
||||
else:
|
||||
return ("Unrecognized or invalid path",)
|
||||
Reference in New Issue
Block a user