42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
from .FL_train_utils import Utils
|
|
import os
|
|
from PIL import Image
|
|
|
|
class FL_LoadImagesFromDirectoryPath:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"directory": ("STRING", {"default": "X://path/to/images"}),
|
|
"caption_extension": ([".caption", ".txt"], {"default": ".txt"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("images", "captions")
|
|
|
|
FUNCTION = "start"
|
|
|
|
CATEGORY = "🏵️Fill Nodes/Training"
|
|
|
|
def start(self, directory, caption_extension):
|
|
|
|
images = []
|
|
captions = []
|
|
if not os.path.exists(directory):
|
|
return (Utils.list_tensor2tensor([]), [])
|
|
|
|
files = Utils.listdir(directory)
|
|
image_files = [f for f in files if f.lower().endswith((".png", ".jpg", ".webp", ".jpeg"))]
|
|
|
|
for image_file in image_files:
|
|
image_path = os.path.join(directory, image_file)
|
|
caption_path = os.path.splitext(image_path)[0] + caption_extension
|
|
|
|
if os.path.exists(caption_path):
|
|
with open(caption_path, 'r', encoding='utf-8') as f:
|
|
captions.append(f.read().strip())
|
|
pil_image = Image.open(image_path)
|
|
images.append(Utils.pil2tensor(pil_image))
|
|
|
|
return (Utils.list_tensor2tensor(images), captions) |