nodes/daxcay-ComfyUI-DataSet
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Major update

Browse Source 
Made changes to make nodes more cohesive
This commit is contained in:
daxcay
2024-07-05 09:04:31 +05:30
parent 37d33ca268
commit e8752c4f71
20 changed files with 441 additions and 543 deletions
Download Patch File Download Diff File Expand all files Collapse all files
classes/DATASET_LoadImage.py
+4 -4
View File
@@ -7,7 +7,7 @@ import hashlib
from PIL import Image, ImageOps, ImageSequence, ImageFile
#code credit: nodes.py comfui
class DATASET_LoadImage:
class DataSet_LoadImage:
@classmethod
def INPUT_TYPES(s):
input_dir = folder_paths.get_input_directory()
@@ -18,7 +18,7 @@ class DATASET_LoadImage:
},
}
CATEGORY = "image"
CATEGORY = "🔶DATASET🔶"
RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING", "STRING", "STRING")
RETURN_NAMES = ("image", "image_mask", "image_name", "image_name_without_extension", "image_path", "image_directory_path")
@@ -89,9 +89,9 @@ class DATASET_LoadImage:
N_CLASS_MAPPINGS = {
"DATASET_LoadImage": DATASET_LoadImage,
"DataSet_LoadImage": DataSet_LoadImage,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_LoadImage": "DATASET_LoadImage",
"DataSet_LoadImage": "DataSet_LoadImage",
}
classes/DATASET_OpenAIChat.py
+3 -8
View File
@@ -1,7 +1,6 @@
from openai import OpenAI
import os
class DATASET_OpenAIChat:
class DataSet_OpenAIChat:
def __init__(self):
pass
@@ -34,14 +33,10 @@ class DATASET_OpenAIChat:
except Exception as e:
return (f"Error: {str(e)}",)
@classmethod
def IS_CHANGED(s, model, api_url, api_key, prompt, token_length):
return os.urandom(16).hex()
N_CLASS_MAPPINGS = {
"DATASET_OpenAIChat": DATASET_OpenAIChat,
"DataSet_OpenAIChat": DataSet_OpenAIChat,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_OpenAIChat": "DATASET_OpenAIChat",
"DataSet_OpenAIChat": "DataSet_OpenAIChat",
}
classes/DATASET_OpenAIChatImage.py
+3 -8
View File
@@ -3,9 +3,8 @@ import io
from PIL import Image
import numpy as np
from openai import OpenAI
import os
class DATASET_OpenAIChatImage:
class DataSet_OpenAIChatImage:
def __init__(self):
pass
@@ -54,14 +53,10 @@ class DATASET_OpenAIChatImage:
except Exception as e:
return (f"Error: {str(e)}",)
@classmethod
def IS_CHANGED(s, image, image_detail, model, api_url, api_key, prompt, token_length):
return os.urandom(16).hex()
N_CLASS_MAPPINGS = {
"DATASET_OpenAIChatImage": DATASET_OpenAIChatImage,
"DataSet_OpenAIChatImage": DataSet_OpenAIChatImage,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_OpenAIChatImage": "DATASET_OpenAIChatImage",
"DataSet_OpenAIChatImage": "DataSet_OpenAIChatImage",
}
classes/DATASET_OpenAIChatImageBatch.py
+3 -7
View File
@@ -5,7 +5,7 @@ import numpy as np
from openai import OpenAI
import os
class DATASET_OpenAIChatImageBatch:
class DataSet_OpenAIChatImageBatch:
def __init__(self):
pass
@@ -70,15 +70,11 @@ class DATASET_OpenAIChatImageBatch:
except Exception as e:
return (f"Error: {str(e)}",)
@classmethod
def IS_CHANGED(s, image, image_detail, model, api_url, api_key, prompt, token_length):
return os.urandom(16).hex()
N_CLASS_MAPPINGS = {
"DATASET_OpenAIChatImageBatch": DATASET_OpenAIChatImageBatch,
"DataSet_OpenAIChatImageBatch": DataSet_OpenAIChatImageBatch,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_OpenAIChatImageBatch": "DATASET_OpenAIChatImageBatch",
"DataSet_OpenAIChatImageBatch": "DataSet_OpenAIChatImageBatch",
}
classes/DATASET_SaveImage.py
+7 -7
View File
@@ -6,7 +6,7 @@ from PIL import Image
from PIL.PngImagePlugin import PngInfo
from comfy.cli_args import args
class DATASET_SaveImage:
class DataSet_SaveImage:
def __init__(self):
self.compression = 4
@@ -16,7 +16,7 @@ class DATASET_SaveImage:
return {
"required": {
"Images": ("IMAGE",),
"Directory": ("STRING", {}),
"destination": ("STRING", {}),
"Filename": ("STRING", {"default": "Image"}),
},
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
@@ -25,13 +25,13 @@ class DATASET_SaveImage:
RETURN_TYPES = ()
FUNCTION = "BatchSave"
OUTPUT_NODE = True
CATEGORY = "🔵 JDCN 🔵"
CATEGORY = "🔶DATASET🔶"
def BatchSave(self, Images, Directory, Filename, prompt=None, extra_pnginfo=None):
def BatchSave(self, Images, destination, Filename, prompt=None, extra_pnginfo=None):
try:
Directory = Directory
Directory = destination
Filename = Filename
if not os.path.exists(Directory):
@@ -61,9 +61,9 @@ class DATASET_SaveImage:
N_CLASS_MAPPINGS = {
"DATASET_SaveImage": DATASET_SaveImage,
"DataSet_SaveImage": DataSet_SaveImage,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_SaveImage": "DATASET_SaveImage",
"DataSet_SaveImage": "DataSet_SaveImage",
}
classes/DATASET_SearchAndReplace.py
-82
View File
@@ -1,82 +0,0 @@
import os
def read_text_file(file_path):
try:
with open(file_path, 'r') as file:
text = file.read()
return text
except FileNotFoundError:
print(f"Error: File '{file_path}' not found.")
return None
except Exception as e:
print(f"Error: {e}")
return None
def save_file(file_path, content):
with open(file_path, 'w') as file:
file.write(content)
print(f"File saved successfully at {file_path}")
class DATASET_SearchAndReplace:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"TextFilePathList": ("STRING", {"forceInput": True}),
"SearchFor": ("STRING", {"multiline": True, "default": "concept"}),
"ReplaceWith": ("STRING", {"multiline": True, "default": "concept"})
},
}
INPUT_IS_LIST = True
RETURN_TYPES = ()
RETURN_NAMES = ()
FUNCTION = "SaR"
OUTPUT_NODE = True
CATEGORY = "🔶DATASET🔶"
def SaR(self, TextFilePathList, SearchFor, ReplaceWith):
try:
directory_path = os.path.dirname(TextFilePathList[0])
backup_path = os.path.join(directory_path, "sr_backup")
os.makedirs(backup_path, exist_ok=True)
for file_path in TextFilePathList:
base_name = os.path.basename(file_path)
save_path = os.path.join(backup_path, base_name)
content = read_text_file(file_path)
save_file(save_path, content)
content = content.replace(SearchFor[0], ReplaceWith[0])
save_file(file_path, content)
return ()
except Exception as e:
print(f"Error saving: {e}")
return ()
@classmethod
def IS_CHANGED(s, TextFilePathList, SearchFor, ReplaceWith):
return os.urandom(16).hex()
N_CLASS_MAPPINGS = {
"DATASET_SearchAndReplace": DATASET_SearchAndReplace,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_SearchAndReplace": "DATASET_SearchAndReplace",
}
classes/DATASET_TXTFileLoader.py
-61
View File
@@ -1,61 +0,0 @@
import os
def read_text_file(file_path):
try:
with open(file_path, 'r') as file:
text = file.read()
return text
except FileNotFoundError:
print(f"Error: File '{file_path}' not found.")
return None
except Exception as e:
print(f"Error: {e}")
return None
class DATASET_TXTFileLoader:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"file_paths": ("STRING", {"forceInput": True}),
},
}
INPUT_IS_LIST = True
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("content",)
OUTPUT_IS_LIST = (True,)
FUNCTION = "LoadIT"
OUTPUT_NODE = True
CATEGORY = "🔶DATASET🔶"
def LoadIT(self, file_paths):
files = []
for file_path in file_paths:
content = read_text_file(file_path)
if(content):
files.append(content)
else:
print(f'Error reading file: {file_path}')
return (files,)
@classmethod
def IS_CHANGED(s, file_paths):
return os.urandom(16).hex()
N_CLASS_MAPPINGS = {
"DATASET_TXTFileLoader": DATASET_TXTFileLoader,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_TXTFileLoader": "DATASET_TXTFileLoader",
}
classes/DATASET_TXTFileSaver.py
-82
View File
@@ -1,82 +0,0 @@
import os
def save_file(filename, output_dir, content, mode='SaveNew'):
os.makedirs(output_dir, exist_ok=True)
file_path = os.path.join(output_dir, filename)
if mode == 'SaveNew':
counter = 0
while os.path.exists(file_path):
counter += 1
file_path = os.path.join(output_dir, f"{os.path.splitext(filename)[0]}_{counter}{os.path.splitext(filename)[1]}")
elif mode == 'Merge' and os.path.exists(file_path):
with open(file_path, 'a') as file:
file.write(content)
print(f"Content appended successfully to {file_path}")
return
elif mode == 'Overwrite' and os.path.exists(file_path):
os.remove(file_path)
elif mode == 'MergeAndSaveNew' and os.path.exists(file_path):
with open(file_path, 'r') as file:
existing_content = file.read()
content = existing_content + content
counter = 0
while os.path.exists(file_path):
counter += 1
file_path = os.path.join(output_dir, f"{os.path.splitext(filename)[0]}_{counter}{os.path.splitext(filename)[1]}")
with open(file_path, 'w') as file:
file.write(content)
print(f"File saved successfully at {file_path}")
class DATASET_TXTFileSaver:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"content": ("STRING",{"forceInput": True}),
"file_name": ("STRING",{"forceInput": True}),
"save_in": ("STRING", {"default": "directory path"}),
"save_mode": (['Merge','Overwrite','SaveNew','MergeAndSaveNew'],),
},
}
# INPUT_IS_LIST = True
RETURN_TYPES = ()
FUNCTION = "SaveIT"
OUTPUT_NODE = True
CATEGORY = "🔶DATASET🔶"
def SaveIT(self, content, file_name, save_in, save_mode):
try:
directory = save_in
mode = save_mode
if not os.path.exists(directory):
os.makedirs(directory)
save_file(f"{file_name}.txt", directory, content, mode)
except Exception as e:
print(f"Error saving: {e}")
return ()
@classmethod
def IS_CHANGED(s, content, file_name, directory, mode):
return os.urandom(16).hex()
N_CLASS_MAPPINGS = {
"DATASET_TXTFileSaver": DATASET_TXTFileSaver,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_TXTFileSaver": "DATASET_TXTFileSaver",
}
classes/DATASET_TagManipulatorByImageNames.py
-150
View File
@@ -1,150 +0,0 @@
import os
def append_text(tags, combined_texts):
tags_list = tags.split(',')
temp_tags = []
for tag in tags_list:
t = tag.strip()
temp_tags.append(t)
tags_list = temp_tags
text_pos_list = [text_pos.strip().split()
for text_pos in combined_texts.split(',')]
for text_pos in text_pos_list:
num_pos = len(text_pos)-1
text = " ".join(text_pos[0:num_pos])
pos = int(text_pos[num_pos])
if pos == 0:
tags_list.insert(0, text)
else:
tags_list.insert(pos, text)
new_tags = ', '.join(tags_list)
return new_tags
def save_file(filename, output_dir, content, mode='Overwrite'):
os.makedirs(output_dir, exist_ok=True)
file_path = os.path.join(output_dir, filename)
if mode == 'SaveNew':
counter = 0
while os.path.exists(file_path):
counter += 1
file_path = os.path.join(
output_dir, f"{os.path.splitext(filename)[0]}_{counter}{os.path.splitext(filename)[1]}")
elif mode == 'Merge' and os.path.exists(file_path):
with open(file_path, 'a') as file:
file.write(content)
print(f"Content appended successfully to {file_path}")
return
elif mode == 'Overwrite' and os.path.exists(file_path):
os.remove(file_path)
elif mode == 'MergeAndSaveNew' and os.path.exists(file_path):
with open(file_path, 'r') as file:
existing_content = file.read()
content = existing_content + content
counter = 0
while os.path.exists(file_path):
counter += 1
file_path = os.path.join(
output_dir, f"{os.path.splitext(filename)[0]}_{counter}{os.path.splitext(filename)[1]}")
with open(file_path, 'w') as file:
file.write(content)
print(f"File saved successfully at {file_path}")
def read_text_file(file_path):
try:
with open(file_path, 'r') as file:
text = file.read()
return text
except FileNotFoundError:
print(f"Error: File '{file_path}' not found.")
return None
except Exception as e:
print(f"Error: {e}")
return None
def read_files_from_directory(file_names, directory):
file_contents = {}
try:
for name in file_names:
file_path = os.path.join(directory, name + ".txt")
if os.path.isfile(file_path):
content = read_text_file(file_path)
if content is not None:
file_contents[name] = content
else:
print(
f"Warning: File '{name}.txt' not found in directory '{directory}'.")
except Exception as e:
print(f"Error: {e}")
return file_contents
class DATASET_TagManipulatorByImageNames:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"ImageNames": ("STRING", {"forceInput": True}),
"CaptionDirectory": ("STRING", {"default": "directory path"}),
"Backup": ("BOOLEAN", {"default":False}),
"Captions": ("STRING", {"multiline": True, "default": "concept"}),
},
}
INPUT_IS_LIST = True
RETURN_TYPES = ()
FUNCTION = "SaveIT"
OUTPUT_NODE = True
CATEGORY = "🔶DATASET🔶"
def SaveIT(self, ImageNames, CaptionDirectory, Captions, Backup):
try:
TagsDirectory = CaptionDirectory
backup_file_path = os.path.join(TagsDirectory[0], "backup")
os.makedirs(backup_file_path, exist_ok=True)
Contents = read_files_from_directory(ImageNames, TagsDirectory[0])
if Backup[0]:
for name in ImageNames:
save_file(f"{name}.txt",backup_file_path,Contents[name],"Overwrite")
for name in ImageNames:
Contents[name] = append_text(Contents[name], Captions[0])
for name in ImageNames:
save_file(f"{name}.txt",TagsDirectory[0],Contents[name],"Overwrite")
except Exception as e:
print(f"Error saving: {e}")
return ()
@classmethod
def IS_CHANGED(s, ImageNames, CaptionDirectory, Captions, Backup):
return os.urandom(16).hex()
N_CLASS_MAPPINGS = {
"DATASET_TagManipulatorByImageNames": DATASET_TagManipulatorByImageNames,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_TagManipulatorByImageNames": "DATASET_TagManipulatorByImageNames",
}
classes/DATASET_TriggerWords.py
+8 -8
View File
@@ -12,7 +12,7 @@ def find_trigger_words(input_str, include_till_comma=False):
return ', '.join(trigger_words)
class DATASET_TriggerWords:
class DataSet_TriggerWords:
def __init__(self):
pass
@@ -21,7 +21,7 @@ class DATASET_TriggerWords:
def INPUT_TYPES(s):
return {
"required": {
"content": ("STRING",{"forceInput": True}),
"TextFileContents": ("STRING",{"forceInput": True}),
"search": (['trigger_word_only','trigger_word_phrase'],),
},
}
@@ -35,17 +35,17 @@ class DATASET_TriggerWords:
CATEGORY = "🔶DATASET🔶"
def FindIT(self, content, search):
def FindIT(self, TextFileContents, search):
try:
search = search[0]
words = []
for sentence in content:
for content in TextFileContents:
if search == "trigger_word_only":
words.append(find_trigger_words(sentence, False))
words.append(find_trigger_words(content, False))
elif search == "trigger_word_phrase":
words.append(find_trigger_words(sentence, True))
words.append(find_trigger_words(content, True))
except Exception as e:
print(f"Error saving: {e}")
@@ -53,9 +53,9 @@ class DATASET_TriggerWords:
return (words,)
N_CLASS_MAPPINGS = {
"DATASET_TriggerWords": DATASET_TriggerWords,
"DataSet_TriggerWords": DataSet_TriggerWords,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_TriggerWords": "DATASET_TriggerWords",
"DataSet_TriggerWords": "DataSet_TriggerWords",
}
classes/DataSet_ConceptManager.py
+72
View File
@@ -0,0 +1,72 @@
def append_text(tags, combined_texts):
tags_list = [tag.strip() for tag in tags.split(',')]
text_pos_list = [text_pos.strip().split() for text_pos in combined_texts.split(',')]
for text_pos in text_pos_list:
num_pos = len(text_pos) - 1
text = " ".join(text_pos[0:num_pos])
pos = int(text_pos[num_pos])
if pos == 0:
tags_list.insert(0, text)
else:
tags_list.insert(pos, text)
new_tags = ', '.join(tags_list)
return new_tags
def remove_text(tags, combined_texts):
tags_list = [tag.strip() for tag in tags.split(',')]
text_pos_list = [text_pos.strip().split() for text_pos in combined_texts.split(',')]
for text_pos in text_pos_list:
num_pos = len(text_pos) - 1
text = " ".join(text_pos[0:num_pos])
if text in tags_list:
tags_list.remove(text)
new_tags = ', '.join(tags_list)
return new_tags
class DataSet_ConceptManager:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"TextFileContents": ("STRING", {"forceInput": True}),
"Mode": (['add','remove'],),
"Concepts": ("STRING", {"multiline": True, "default": "concept"}),
},
}
INPUT_IS_LIST = True
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("TextFileContents")
OUTPUT_IS_LIST = (True,)
FUNCTION = "SaveIT"
OUTPUT_NODE = True
CATEGORY = "🔶DATASET🔶"
def SaveIT(self, TextFileContents, Mode, Concepts):
try:
edited = []
for content in TextFileContents:
if Mode[0] == "add":
edited.append(append_text(content, Concepts[0]))
elif Mode[0] == "remove":
edited.append(remove_text(content, Concepts[0]))
except Exception as e:
print(f"Error saving: {e}")
return (edited)
N_CLASS_MAPPINGS = {
"DataSet_ConceptManager": DataSet_ConceptManager,
}
N_DISPLAY_NAME_MAPPINGS = {
"DataSet_ConceptManager": "DataSet_ConceptManager",
}
classes/DATASET_xCopy.py → classes/DataSet_CopyFiles.py
+14 -9
View File
@@ -5,8 +5,13 @@ def blind_copy(source, destination):
os.makedirs(destination, exist_ok=True)
for filename in os.listdir(source):
source_file = os.path.join(source, filename)
dest_file = os.path.join(destination, filename)
if os.path.isfile(source_file):
shutil.copy(source_file, destination)
try:
shutil.copy(source_file, dest_file)
print(f"Copied {source_file} to {dest_file}")
except Exception as e:
print(f"Failed to copy {source_file} to {dest_file}: {e}")
def copy_by_destination_files(source, destination):
os.makedirs(destination, exist_ok=True)
@@ -17,9 +22,13 @@ def copy_by_destination_files(source, destination):
file_base_name = os.path.splitext(filename)[0]
dest_file_match = any(os.path.splitext(dest_filename)[0] == file_base_name for dest_filename in dest_files)
if dest_file_match and filename not in dest_files:
shutil.copy(source_file, destination)
try:
shutil.copy(source_file, destination)
print(f"Copied {source_file} to {destination}")
except Exception as e:
print(f"Failed to copy {source_file} to {destination}: {e}")
class DATASET_xCopy:
class DataSet_CopyFiles:
def __init__(self):
pass
@@ -56,14 +65,10 @@ class DATASET_xCopy:
return ()
@classmethod
def IS_CHANGED(s, source_folder, destination_folder, mode):
return os.urandom(16).hex()
N_CLASS_MAPPINGS = {
"DATASET_xCopy": DATASET_xCopy,
"DataSet_CopyFiles": DataSet_CopyFiles,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_xCopy": "DATASET_xCopy",
"DataSet_CopyFiles": "DataSet_CopyFiles",
}
classes/DataSet_FindAndReplace.py
+36
View File
@@ -0,0 +1,36 @@
class DataSet_FindAndReplace:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"TextFileContents": ("STRING", {"forceInput": True}),
"SearchFor": ("STRING", {"multiline": True, "default": "concept"}),
"ReplaceWith": ("STRING", {"multiline": True, "default": "concept"})
},
}
INPUT_IS_LIST = True
RETURN_TYPES = ("STRING", )
RETURN_NAMES = ("TextFileContents",)
OUTPUT_IS_LIST = (True,)
FUNCTION = "SaR"
OUTPUT_NODE = True
CATEGORY = "🔶DATASET🔶"
def SaR(self, TextFileContents, SearchFor, ReplaceWith):
edited = []
for content in TextFileContents:
edited.append(content.replace(SearchFor[0], ReplaceWith[0]))
return (edited,)
N_CLASS_MAPPINGS = {
"DataSet_FindAndReplace": DataSet_FindAndReplace,
}
N_DISPLAY_NAME_MAPPINGS = {
"DataSet_FindAndReplace": "DataSet_FindAndReplace",
}
classes/DataSet_PathSelector.py
+70
View File
@@ -0,0 +1,70 @@
import os
def search_and_select_files(search_in_directory, search_for_extension, select_from_directory, select_extension):
# Search for files with the specified extension in the search_in_directory
search_results = [file for file in os.listdir(search_in_directory) if file.endswith(search_for_extension)]
search_results_names = {os.path.splitext(file)[0] for file in search_results}
# List all files with the specified extension in the select_from_directory
selection_files = [file for file in os.listdir(select_from_directory) if file.endswith(select_extension)]
selection_files_dict = {os.path.splitext(file)[0]: file for file in selection_files}
# Find the matching files based on the names
matching_files = [selection_files_dict[name] for name in search_results_names if name in selection_files_dict]
# Prepare the return values
matching_files_full_paths = [os.path.join(select_from_directory, file) for file in matching_files]
matching_files_with_extensions = matching_files
matching_files_without_extensions = [os.path.splitext(file)[0] for file in matching_files]
return (matching_files_with_extensions, matching_files_without_extensions, matching_files_full_paths)
class DataSet_PathSelector:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"search_in_directory": ("STRING", {"multiline": False, "default": ""}),
"search_for_extension": ("STRING", {"multiline": False, "default": ""}),
"select_from_directory": ("STRING", {"multiline": False, "default": ""}),
"select_extension": ("STRING", {"multiline": False, "default": ""}),
},
}
RETURN_TYPES = ("STRING", "STRING", "INT",)
RETURN_NAMES = ("SelectedNamesWithExtension","SelectedNamesWithoutExtension", "SelectedPaths")
OUTPUT_IS_LIST = (True, True, True)
OUTPUT_NODE = True
FUNCTION = "make_list"
CATEGORY = "🔶DATASET🔶"
def make_list(self, search_in_directory, search_for_extension, select_from_directory, select_extension):
try:
if not os.path.exists(search_in_directory):
print(f"The folder '{search_in_directory}' does not exist.")
return ([], [], 0)
if not os.path.exists(select_from_directory):
print(f"The folder '{select_from_directory}' does not exist.")
return ([], [], 0)
search_for_extension = search_for_extension if search_for_extension.startswith('.') else '.' + search_for_extension
select_extension = select_extension if select_extension.startswith('.') else '.' + select_extension
a,b,c = search_and_select_files(search_in_directory, search_for_extension, select_from_directory, select_extension)
return (a,b,c)
except Exception as e:
print(f"An error occurred: {e}")
return ([], [], [])
N_CLASS_MAPPINGS = {
"DataSet_PathSelector": DataSet_PathSelector,
}
N_DISPLAY_NAME_MAPPINGS = {
"DataSet_PathSelector": "DataSet_PathSelector",
}
classes/DataSet_TextFilesLoad.py
+67
View File
@@ -0,0 +1,67 @@
import os
def read_text_file(file_path):
try:
with open(file_path, 'r') as file:
text = file.read()
return text
except FileNotFoundError:
print(f"Error: File '{file_path}' not found.")
return None
except Exception as e:
print(f"Error: {e}")
return None
class DataSet_TextFilesLoad:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"directory": ("STRING", {"default":"directory_path"}),
},
}
INPUT_IS_LIST = True
RETURN_TYPES = ("STRING","STRING","STRING",)
RETURN_NAMES = ("TextFileNames","TextFileNamesWithoutExtension","TextFilePaths","TextFileContents",)
OUTPUT_IS_LIST = (True, True, True, True,)
FUNCTION = "LoadIT"
OUTPUT_NODE = True
CATEGORY = "🔶DATASET🔶"
def LoadIT(self, directory):
try:
directory = directory[0]
file_paths = [os.path.join(directory, f) for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f)) and f.endswith('.txt')]
file_names = [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f)) and f.endswith('.txt')]
file_names_without_extension = [os.path.splitext(f)[0] for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f)) and f.endswith('.txt')]
file_contents = []
for file_path in file_paths:
content = read_text_file(file_path)
if(content):
file_contents.append(content)
else:
print(f'Error reading file: {file_path}')
return (file_names, file_names_without_extension, file_paths, file_contents,)
except Exception as e:
return (file_names, file_names_without_extension, file_paths, file_contents,)
N_CLASS_MAPPINGS = {
"DataSet_TextFilesLoad": DataSet_TextFilesLoad,
}
N_DISPLAY_NAME_MAPPINGS = {
"DataSet_TextFilesLoad": "DataSet_TextFilesLoad",
}
classes/DataSet_TextFilesLoadFromList.py
+65
View File
@@ -0,0 +1,65 @@
import os
def read_text_file(file_path):
try:
with open(file_path, 'r') as file:
text = file.read()
return text
except FileNotFoundError:
print(f"Error: File '{file_path}' not found.")
return None
except Exception as e:
print(f"Error: {e}")
return None
class DataSet_TextFilesLoadFromList:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"TextFilePathsList": ("STRING", {"forceInput":True}),
},
}
INPUT_IS_LIST = True
RETURN_TYPES = ("STRING","STRING","STRING",)
RETURN_NAMES = ("TextFileNames","TextFileNamesWithoutExtension","TextFilePaths","TextFileContents",)
OUTPUT_IS_LIST = (True, True, True, True,)
FUNCTION = "LoadIT"
OUTPUT_NODE = True
CATEGORY = "🔶DATASET🔶"
def LoadIT(self, TextFilePathsList):
try:
file_paths = [path for path in TextFilePathsList if path.endswith('.txt')]
file_names = [os.path.basename(path) for path in file_paths]
file_names_without_ext = [os.path.splitext(os.path.basename(path))[0] for path in file_paths]
file_contents = []
for file_path in file_paths:
content = read_text_file(file_path)
if(content):
file_contents.append(content)
else:
print(f'Error reading file: {file_path}')
return (file_names, file_names_without_ext, file_paths, file_contents,)
except Exception as e:
return ([],[],[])
N_CLASS_MAPPINGS = {
"DataSet_TextFilesLoadFromList": DataSet_TextFilesLoadFromList,
}
N_DISPLAY_NAME_MAPPINGS = {
"DataSet_TextFilesLoadFromList": "DataSet_TextFilesLoadFromList",
}
classes/DATASET_TXTFileSaverBatch.py → classes/DataSet_TextFilesSave.py
+11 -16
View File
@@ -29,7 +29,7 @@ def save_file(filename, output_dir, content, mode='SaveNew'):
file.write(content)
print(f"File saved successfully at {file_path}")
class DATASET_TXTFileSaverBatch:
class DataSet_TextFilesSave:
def __init__(self):
pass
@@ -38,9 +38,9 @@ class DATASET_TXTFileSaverBatch:
def INPUT_TYPES(s):
return {
"required": {
"file_names": ("STRING",{"forceInput": True}),
"contents": ("STRING",{"forceInput": True}),
"save_in": ("STRING", {"default": "directory path"}),
"TextFileContents": ("STRING",{"forceInput": True}),
"TextFileNames": ("STRING",{"forceInput": True}),
"destination": ("STRING", {"default": "directory path"}),
"save_mode": (['Merge','Overwrite','SaveNew','MergeAndSaveNew'],),
},
}
@@ -52,34 +52,29 @@ class DATASET_TXTFileSaverBatch:
CATEGORY = "🔶DATASET🔶"
def SaveIT(self, file_names, contents, save_in, save_mode):
def SaveIT(self, TextFileContens, TextFileNames, destination, save_mode):
try:
directory = save_in[0]
directory = destination[0]
mode = save_mode[0]
if not os.path.exists(directory):
os.makedirs(directory)
for i in range(0, len(contents)):
text = contents[i]
file_name = file_names[i]
for i in range(0, len(TextFileContens)):
text = TextFileContens[i]
file_name = TextFileNames[i]
save_file(f"{file_name}.txt", directory, text, mode)
except Exception as e:
print(f"Error saving: {e}")
return ()
@classmethod
def IS_CHANGED(s, content, file_name, directory, mode):
return os.urandom(16).hex()
N_CLASS_MAPPINGS = {
"DATASET_TXTFileSaverBatch": DATASET_TXTFileSaverBatch,
"DataSet_TextFilesSave": DataSet_TextFilesSave,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_TXTFileSaverBatch": "DATASET_TXTFileSaverBatch",
"DataSet_TextFilesSave": "DataSet_TextFilesSave",
}
classes/DATASET_CaptionVisualizer.py → classes/DataSet_Visualizer.py
+40 -65
View File
@@ -10,25 +10,29 @@ import networkx as nx
from collections import Counter, defaultdict
from itertools import combinations
import pandas as pd
import folder_paths
def generate_wordcloud_and_network_graph(contents, separator, top_n_wordcloud=100, top_n_network=100, top_n_table=10):
output_dir = folder_paths.get_output_directory()
separators = {"comma":',','colon':';','space':' ','pipe':'|'}
separator = separators.get(separator, ',')
def generate_wordcloud_and_network_graph(file_paths, output_dir, top_n_wordcloud=100, top_n_network=100, top_n_table=10):
word_counter = Counter()
tag_cooccurrences = Counter()
for file_path in file_paths:
for content in contents:
try:
with open(file_path, 'r', encoding='utf-8') as file:
tags = file.read().strip().split(',')
word_counter.update(tags)
for tag1 in tags:
for tag2 in tags:
if tag1 != tag2:
tag_cooccurrences[(tag1, tag2)] += 1
tags = content.strip().split(separator)
word_counter.update(tags)
for tag1 in tags:
for tag2 in tags:
if tag1 != tag2:
tag_cooccurrences[(tag1, tag2)] += 1
except Exception as e:
print(f"Error reading file {file_path}: {e}")
print(f"Error processing content: {e}")
# Word cloud generation
top_wordcloud_tags = dict(word_counter.most_common(top_n_wordcloud))
wordcloud = WordCloud(width=1920, height=1080, background_color='white').generate_from_frequencies(top_wordcloud_tags)
@@ -42,24 +46,19 @@ def generate_wordcloud_and_network_graph(file_paths, output_dir, top_n_wordcloud
print("Word cloud saved as", output_wordcloud_file)
# Network graph generation
G = nx.Graph()
tags_cooccurrence = defaultdict(int)
for file_path in file_paths:
for content in contents:
try:
with open(file_path, 'r', encoding='utf-8') as file:
content = file.read()
tags = list(set(content.split(',')))
for tag_pair in combinations(tags, 2):
if tag_pair[0].strip() and tag_pair[1].strip():
tags_cooccurrence[tag_pair] += 1
tags = list(set(content.split(separator)))
for tag_pair in combinations(tags, 2):
if tag_pair[0].strip() and tag_pair[1].strip():
tags_cooccurrence[tag_pair] += 1
except Exception as e:
print(f"Error reading file {file_path}: {e}")
print(f"Error processing content: {e}")
top_cooccurrences = sorted(tags_cooccurrence.items(
), key=lambda x: x[1], reverse=True)[:top_n_network]
top_cooccurrences = sorted(tags_cooccurrence.items(), key=lambda x: x[1], reverse=True)[:top_n_network]
for (tag1, tag2), weight in top_cooccurrences:
G.add_edge(tag1.strip(), tag2.strip(), weight=weight)
@@ -74,61 +73,44 @@ def generate_wordcloud_and_network_graph(file_paths, output_dir, top_n_wordcloud
edge_width = [G[u][v]['weight'] / 100 for u, v in G.edges]
pos = nx.kamada_kawai_layout(G)
nx.draw_networkx_nodes(G, pos, node_size=node_size,
node_color=node_color, cmap=plt.cm.plasma, alpha=0.8)
nx.draw_networkx_nodes(G, pos, node_size=node_size, node_color=node_color, cmap=plt.cm.plasma, alpha=0.8)
nx.draw_networkx_edges(G, pos, width=edge_width, alpha=0.3, edge_color='w')
nx.draw_networkx_labels(G, pos, font_size=12,
font_weight='bold', font_color='white',
font_family='sans-serif')
nx.draw_networkx_labels(G, pos, font_size=12, font_weight='bold', font_color='white', font_family='sans-serif')
plt.axis('off')
output_network_graph_file = join(output_dir, 'network_graph.png')
plt.savefig(output_network_graph_file, pad_inches=0, dpi=300,
bbox_inches='tight', facecolor=gradio_blue)
plt.savefig(output_network_graph_file, pad_inches=0, dpi=300, bbox_inches='tight', facecolor=gradio_blue)
plt.close()
print("Network graph saved as", output_network_graph_file)
# Tag frequency table
tag_freq_table = pd.DataFrame.from_dict(
word_counter, orient='index', columns=['Frequency'])
tag_freq_table = tag_freq_table.sort_values(
by='Frequency', ascending=False)
tag_freq_table = pd.DataFrame.from_dict(word_counter, orient='index', columns=['Frequency'])
tag_freq_table = tag_freq_table.sort_values(by='Frequency', ascending=False)
tag_freq_table.reset_index(inplace=True)
tag_freq_table.columns = ['Tag', 'Frequency']
plt.figure(figsize=(12, 6))
bars = plt.bar(tag_freq_table['Tag'][:top_n_table],
tag_freq_table['Frequency'][:top_n_table], color='skyblue')
bars = plt.bar(tag_freq_table['Tag'][:top_n_table], tag_freq_table['Frequency'][:top_n_table], color='skyblue')
for bar in bars:
yval = bar.get_height()
plt.text(bar.get_x() + bar.get_width()/2, yval,
round(yval, 2), ha='center', va='bottom', color='black')
plt.text(bar.get_x() + bar.get_width() / 2, yval, round(yval, 2), ha='center', va='bottom', color='black')
plt.xlabel('Tag')
plt.ylabel('Frequency')
plt.xticks(rotation=45, ha='right')
plt.tight_layout()
plt.savefig(join(output_dir, 'tag_frequency_table.png'))
output_table_file = join(output_dir, 'tag_frequency_table.png')
plt.savefig(output_table_file)
plt.close()
# plt.figure(figsize=(12, 6))
# plt.bar(tag_freq_table['Tag'][:top_n_table], tag_freq_table['Frequency'][:top_n_table], color='skyblue')
# plt.xlabel('Tag')
# plt.ylabel('Frequency')
# plt.xticks(rotation=45, ha='right')
# plt.tight_layout()
# plt.savefig(join(output_dir, 'tag_frequency_table.png'))
# plt.close()
print("Tag frequency table saved as", output_table_file)
print("Tag frequency table saved as", join(
output_dir, 'tag_frequency_table.png'))
return output_wordcloud_file, output_network_graph_file, output_table_file
return output_wordcloud_file, output_network_graph_file, join(output_dir, 'tag_frequency_table.png')
def pilToImage(image):
@@ -157,7 +139,7 @@ def create_empty_image(width=100, height=100, color=(255, 255, 255)):
return empty_image
class DATASET_CaptionVisualizer:
class DataSet_Visualizer:
def __init__(self):
pass
@@ -166,7 +148,8 @@ class DATASET_CaptionVisualizer:
def INPUT_TYPES(s):
return {
"required": {
"Captions": ("STRING", {"forceInput": True}),
"TextFileContents": ("STRING", {"forceInput": True}),
"Seperator": (['comma', 'colon', 'space', 'pipe'],),
"WordCloudTop": ("INT", {"default": 1, "min": 1, "max": 9999}),
"NetworkGraphTop": ("INT", {"default": 1, "min": 1, "max": 9999}),
"FrequencyGraphTop": ("INT", {"default": 1, "min": 1, "max": 9999})
@@ -181,14 +164,11 @@ class DATASET_CaptionVisualizer:
OUTPUT_NODE = True
CATEGORY = "🔶DATASET🔶"
def Visualize(self, Captions, WordCloudTop, NetworkGraphTop, FrequencyGraphTop):
def Visualize(self, TextFileContents, Seperator, WordCloudTop, NetworkGraphTop, FrequencyGraphTop):
try:
directory_path = os.path.dirname(Captions[0])
visualize_path = os.path.join(directory_path, "visualize")
os.makedirs(visualize_path, exist_ok=True)
wc, ng, fg = generate_wordcloud_and_network_graph(Captions, visualize_path, WordCloudTop[0], NetworkGraphTop[0], FrequencyGraphTop[0])
wc, ng, fg = generate_wordcloud_and_network_graph(TextFileContents, Seperator[0], WordCloudTop[0], NetworkGraphTop[0], FrequencyGraphTop[0])
images = []
if os.path.exists(wc):
@@ -204,16 +184,11 @@ class DATASET_CaptionVisualizer:
print(f"Error saving: {e}")
return (["", "", ""],)
@classmethod
def IS_CHANGED(s, Captions, WordCloudTop, NetworkGraphTop, FrequencyGraphTop):
return os.urandom(16).hex()
N_CLASS_MAPPINGS = {
"DATASET_CaptionVisualizer": DATASET_CaptionVisualizer,
"DataSet_Visualizer": DataSet_Visualizer,
}
N_DISPLAY_NAME_MAPPINGS = {
"DATASET_CaptionVisualizer": "DATASET_CaptionVisualizer",
"DataSet_Visualizer": "DataSet_Visualizer",
}