From e8752c4f719a3a8801af0fdcca94b2edc21fb366 Mon Sep 17 00:00:00 2001 From: daxcay Date: Fri, 5 Jul 2024 09:04:31 +0530 Subject: [PATCH] Major update Made changes to make nodes more cohesive --- __init__.py | 72 +++++---- classes/DATASET_LoadImage.py | 8 +- classes/DATASET_OpenAIChat.py | 11 +- classes/DATASET_OpenAIChatImage.py | 11 +- classes/DATASET_OpenAIChatImageBatch.py | 10 +- classes/DATASET_SaveImage.py | 14 +- classes/DATASET_SearchAndReplace.py | 82 ---------- classes/DATASET_TXTFileLoader.py | 61 ------- classes/DATASET_TXTFileSaver.py | 82 ---------- classes/DATASET_TagManipulatorByImageNames.py | 150 ------------------ classes/DATASET_TriggerWords.py | 16 +- classes/DataSet_ConceptManager.py | 72 +++++++++ ...{DATASET_xCopy.py => DataSet_CopyFiles.py} | 23 +-- classes/DataSet_FindAndReplace.py | 36 +++++ classes/DataSet_PathSelector.py | 70 ++++++++ classes/DataSet_TextFilesLoad.py | 67 ++++++++ classes/DataSet_TextFilesLoadFromList.py | 65 ++++++++ ...SaverBatch.py => DataSet_TextFilesSave.py} | 27 ++-- ...ionVisualizer.py => DataSet_Visualizer.py} | 105 +++++------- web/util.js | 2 +- 20 files changed, 441 insertions(+), 543 deletions(-) delete mode 100644 classes/DATASET_SearchAndReplace.py delete mode 100644 classes/DATASET_TXTFileLoader.py delete mode 100644 classes/DATASET_TXTFileSaver.py delete mode 100644 classes/DATASET_TagManipulatorByImageNames.py create mode 100644 classes/DataSet_ConceptManager.py rename classes/{DATASET_xCopy.py => DataSet_CopyFiles.py} (72%) create mode 100644 classes/DataSet_FindAndReplace.py create mode 100644 classes/DataSet_PathSelector.py create mode 100644 classes/DataSet_TextFilesLoad.py create mode 100644 classes/DataSet_TextFilesLoadFromList.py rename classes/{DATASET_TXTFileSaverBatch.py => DataSet_TextFilesSave.py} (73%) rename classes/{DATASET_CaptionVisualizer.py => DataSet_Visualizer.py} (53%) diff --git a/__init__.py b/__init__.py index 83df9f9..d66c592 100644 --- a/__init__.py +++ b/__init__.py @@ -10,56 +10,58 @@ # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN # THE SOFTWARE. - """ @author: Daxton Caylor @title: ComfyUI-DataSet @nickname: ComfyUI-DataSet -@description: Data Research, Preparation And Manipulators Nodes for Model Trainers, Artists, Designers and Animators. +@description: Advanced Training Data Research, Preparation, and Manipulation Nodes for Expert Model Trainers, Artists, Designers, and Animators. """ -from .classes.DATASET_TXTFileSaver import N_CLASS_MAPPINGS as TXTFileSaverMappings, N_DISPLAY_NAME_MAPPINGS as TXTFileSaverNameMappings -from .classes.DATASET_TXTFileLoader import N_CLASS_MAPPINGS as TXTFileLoaderMappings, N_DISPLAY_NAME_MAPPINGS as TXTFileLoaderNameMappings -from .classes.DATASET_TagManipulatorByImageNames import N_CLASS_MAPPINGS as TagManipulatorByImageNamesMappings, N_DISPLAY_NAME_MAPPINGS as TagManipulatorByImageNamesNameMappings -from .classes.DATASET_CaptionVisualizer import N_CLASS_MAPPINGS as CaptionVisualizerMappings, N_DISPLAY_NAME_MAPPINGS as CaptionVisualizerNameMappings -from .classes.DATASET_SearchAndReplace import N_CLASS_MAPPINGS as SearchAndReplaceMappings, N_DISPLAY_NAME_MAPPINGS as SearchAndReplaceNameMappings -from .classes.DATASET_xCopy import N_CLASS_MAPPINGS as xCopyMappings, N_DISPLAY_NAME_MAPPINGS as xCopyNameMappings -from .classes.DATASET_OpenAIChat import N_CLASS_MAPPINGS as OpenAIChatMappings, N_DISPLAY_NAME_MAPPINGS as OpenAIChatNameMappings -from .classes.DATASET_OpenAIChatImage import N_CLASS_MAPPINGS as OpenAIChatImageMappings, N_DISPLAY_NAME_MAPPINGS as OpenAIChatImageNameMappings -from .classes.DATASET_LoadImage import N_CLASS_MAPPINGS as LoadImageMappings, N_DISPLAY_NAME_MAPPINGS as LoadImageNameMappings -from .classes.DATASET_SaveImage import N_CLASS_MAPPINGS as SaveImageMappings, N_DISPLAY_NAME_MAPPINGS as SaveImageNameMappings -from .classes.DATASET_TXTFileSaverBatch import N_CLASS_MAPPINGS as TXTFileSaverBatchMappings, N_DISPLAY_NAME_MAPPINGS as TXTFileSaverBatchNameMappings -from .classes.DATASET_OpenAIChatImageBatch import N_CLASS_MAPPINGS as OpenAIChatImageBatchMappings, N_DISPLAY_NAME_MAPPINGS as OpenAIChatImageBatchNameMappings -from .classes.DATASET_TriggerWords import N_CLASS_MAPPINGS as TriggerWordsMappings, N_DISPLAY_NAME_MAPPINGS as TriggerWordsNameMappings +from .classes.DataSet_Visualizer import N_CLASS_MAPPINGS as VisualizerMappings, N_DISPLAY_NAME_MAPPINGS as VisualizerNameMappings +from .classes.DataSet_TextFilesLoad import N_CLASS_MAPPINGS as TextFilesLoadMappings, N_DISPLAY_NAME_MAPPINGS as TextFilesLoadNameMappings +from .classes.DataSet_TextFilesLoadFromList import N_CLASS_MAPPINGS as TextFilesLoadFromListMappings, N_DISPLAY_NAME_MAPPINGS as TextFilesLoadFromListNameMappings +from .classes.DataSet_CopyFiles import N_CLASS_MAPPINGS as CopyFilesMappings, N_DISPLAY_NAME_MAPPINGS as CopyFilesNameMappings +from .classes.DataSet_TriggerWords import N_CLASS_MAPPINGS as TriggerWordsMappings, N_DISPLAY_NAME_MAPPINGS as TriggerWordsNameMappings +from .classes.DataSet_TextFilesSave import N_CLASS_MAPPINGS as TextFilesSaveMappings, N_DISPLAY_NAME_MAPPINGS as TextFilesSaveNameMappings +from .classes.DataSet_ConceptManager import N_CLASS_MAPPINGS as ConceptManagerMappings, N_DISPLAY_NAME_MAPPINGS as ConceptManagerNameMappings +from .classes.DataSet_FindAndReplace import N_CLASS_MAPPINGS as FindAndReplaceMappings, N_DISPLAY_NAME_MAPPINGS as FindAndReplaceNameMappings +from .classes.DataSet_PathSelector import N_CLASS_MAPPINGS as PathSelectorMappings, N_DISPLAY_NAME_MAPPINGS as PathSelectorNameMappings +from .classes.DataSet_LoadImage import N_CLASS_MAPPINGS as LoadImageMappings, N_DISPLAY_NAME_MAPPINGS as LoadImageNameMappings +from .classes.DataSet_SaveImage import N_CLASS_MAPPINGS as SaveImageMappings, N_DISPLAY_NAME_MAPPINGS as SaveImageNameMappings +from .classes.DataSet_OpenAIChat import N_CLASS_MAPPINGS as OpenAIChatMappings, N_DISPLAY_NAME_MAPPINGS as OpenAIChatNameMappings +from .classes.DataSet_OpenAIChatImage import N_CLASS_MAPPINGS as OpenAIChatImageMappings, N_DISPLAY_NAME_MAPPINGS as OpenAIChatImageNameMappings +from .classes.DataSet_OpenAIChatImageBatch import N_CLASS_MAPPINGS as OpenAIChatImageBatchMappings, N_DISPLAY_NAME_MAPPINGS as OpenAIChatImageBatchNameMappings NODE_CLASS_MAPPINGS = {} -NODE_CLASS_MAPPINGS.update(TXTFileSaverMappings) -NODE_CLASS_MAPPINGS.update(TXTFileLoaderMappings) -NODE_CLASS_MAPPINGS.update(TagManipulatorByImageNamesMappings) -NODE_CLASS_MAPPINGS.update(CaptionVisualizerMappings) -NODE_CLASS_MAPPINGS.update(SearchAndReplaceMappings) -NODE_CLASS_MAPPINGS.update(xCopyMappings) -NODE_CLASS_MAPPINGS.update(OpenAIChatMappings) -NODE_CLASS_MAPPINGS.update(OpenAIChatImageMappings) +NODE_CLASS_MAPPINGS.update(VisualizerMappings) +NODE_CLASS_MAPPINGS.update(CopyFilesMappings) +NODE_CLASS_MAPPINGS.update(TriggerWordsMappings) +NODE_CLASS_MAPPINGS.update(TextFilesLoadFromListMappings) +NODE_CLASS_MAPPINGS.update(TextFilesLoadMappings) +NODE_CLASS_MAPPINGS.update(TextFilesSaveMappings) +NODE_CLASS_MAPPINGS.update(FindAndReplaceMappings) +NODE_CLASS_MAPPINGS.update(PathSelectorMappings) +NODE_CLASS_MAPPINGS.update(ConceptManagerMappings) NODE_CLASS_MAPPINGS.update(LoadImageMappings) NODE_CLASS_MAPPINGS.update(SaveImageMappings) -NODE_CLASS_MAPPINGS.update(TXTFileSaverBatchMappings) +NODE_CLASS_MAPPINGS.update(OpenAIChatMappings) +NODE_CLASS_MAPPINGS.update(OpenAIChatImageMappings) NODE_CLASS_MAPPINGS.update(OpenAIChatImageBatchMappings) -NODE_CLASS_MAPPINGS.update(TriggerWordsMappings) NODE_DISPLAY_NAME_MAPPINGS = {} -NODE_DISPLAY_NAME_MAPPINGS.update(TXTFileSaverNameMappings) -NODE_DISPLAY_NAME_MAPPINGS.update(TXTFileLoaderNameMappings) -NODE_DISPLAY_NAME_MAPPINGS.update(TagManipulatorByImageNamesNameMappings) -NODE_DISPLAY_NAME_MAPPINGS.update(CaptionVisualizerNameMappings) -NODE_DISPLAY_NAME_MAPPINGS.update(SearchAndReplaceNameMappings) -NODE_DISPLAY_NAME_MAPPINGS.update(xCopyNameMappings) -NODE_DISPLAY_NAME_MAPPINGS.update(OpenAIChatNameMappings) -NODE_DISPLAY_NAME_MAPPINGS.update(OpenAIChatImageNameMappings) +NODE_DISPLAY_NAME_MAPPINGS.update(VisualizerNameMappings) +NODE_DISPLAY_NAME_MAPPINGS.update(CopyFilesNameMappings) +NODE_DISPLAY_NAME_MAPPINGS.update(TriggerWordsNameMappings) +NODE_DISPLAY_NAME_MAPPINGS.update(TextFilesLoadNameMappings) +NODE_DISPLAY_NAME_MAPPINGS.update(TextFilesLoadFromListNameMappings) +NODE_DISPLAY_NAME_MAPPINGS.update(TextFilesSaveNameMappings) +NODE_DISPLAY_NAME_MAPPINGS.update(FindAndReplaceNameMappings) +NODE_DISPLAY_NAME_MAPPINGS.update(PathSelectorNameMappings) +NODE_DISPLAY_NAME_MAPPINGS.update(ConceptManagerNameMappings) NODE_DISPLAY_NAME_MAPPINGS.update(LoadImageNameMappings) NODE_DISPLAY_NAME_MAPPINGS.update(SaveImageNameMappings) -NODE_DISPLAY_NAME_MAPPINGS.update(TXTFileSaverBatchNameMappings) +NODE_DISPLAY_NAME_MAPPINGS.update(OpenAIChatNameMappings) +NODE_DISPLAY_NAME_MAPPINGS.update(OpenAIChatImageNameMappings) NODE_DISPLAY_NAME_MAPPINGS.update(OpenAIChatImageBatchNameMappings) -NODE_DISPLAY_NAME_MAPPINGS.update(TriggerWordsNameMappings) WEB_DIRECTORY = "./web" \ No newline at end of file diff --git a/classes/DATASET_LoadImage.py b/classes/DATASET_LoadImage.py index cb82281..04c84fb 100644 --- a/classes/DATASET_LoadImage.py +++ b/classes/DATASET_LoadImage.py @@ -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", } diff --git a/classes/DATASET_OpenAIChat.py b/classes/DATASET_OpenAIChat.py index 3b8ca21..3324305 100644 --- a/classes/DATASET_OpenAIChat.py +++ b/classes/DATASET_OpenAIChat.py @@ -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", } \ No newline at end of file diff --git a/classes/DATASET_OpenAIChatImage.py b/classes/DATASET_OpenAIChatImage.py index cff885a..8e0fde0 100644 --- a/classes/DATASET_OpenAIChatImage.py +++ b/classes/DATASET_OpenAIChatImage.py @@ -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", } \ No newline at end of file diff --git a/classes/DATASET_OpenAIChatImageBatch.py b/classes/DATASET_OpenAIChatImageBatch.py index f6e0688..5d357d4 100644 --- a/classes/DATASET_OpenAIChatImageBatch.py +++ b/classes/DATASET_OpenAIChatImageBatch.py @@ -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", } \ No newline at end of file diff --git a/classes/DATASET_SaveImage.py b/classes/DATASET_SaveImage.py index 1371a91..1653bc4 100644 --- a/classes/DATASET_SaveImage.py +++ b/classes/DATASET_SaveImage.py @@ -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", } diff --git a/classes/DATASET_SearchAndReplace.py b/classes/DATASET_SearchAndReplace.py deleted file mode 100644 index 6cafad3..0000000 --- a/classes/DATASET_SearchAndReplace.py +++ /dev/null @@ -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", -} diff --git a/classes/DATASET_TXTFileLoader.py b/classes/DATASET_TXTFileLoader.py deleted file mode 100644 index 52b7fc7..0000000 --- a/classes/DATASET_TXTFileLoader.py +++ /dev/null @@ -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", -} diff --git a/classes/DATASET_TXTFileSaver.py b/classes/DATASET_TXTFileSaver.py deleted file mode 100644 index 713c396..0000000 --- a/classes/DATASET_TXTFileSaver.py +++ /dev/null @@ -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", -} diff --git a/classes/DATASET_TagManipulatorByImageNames.py b/classes/DATASET_TagManipulatorByImageNames.py deleted file mode 100644 index ccc4b5b..0000000 --- a/classes/DATASET_TagManipulatorByImageNames.py +++ /dev/null @@ -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", -} diff --git a/classes/DATASET_TriggerWords.py b/classes/DATASET_TriggerWords.py index 3611f30..2b7d3cf 100644 --- a/classes/DATASET_TriggerWords.py +++ b/classes/DATASET_TriggerWords.py @@ -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", } diff --git a/classes/DataSet_ConceptManager.py b/classes/DataSet_ConceptManager.py new file mode 100644 index 0000000..6ee87a9 --- /dev/null +++ b/classes/DataSet_ConceptManager.py @@ -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", +} diff --git a/classes/DATASET_xCopy.py b/classes/DataSet_CopyFiles.py similarity index 72% rename from classes/DATASET_xCopy.py rename to classes/DataSet_CopyFiles.py index 9507a40..d5d8a9c 100644 --- a/classes/DATASET_xCopy.py +++ b/classes/DataSet_CopyFiles.py @@ -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", } diff --git a/classes/DataSet_FindAndReplace.py b/classes/DataSet_FindAndReplace.py new file mode 100644 index 0000000..e026c62 --- /dev/null +++ b/classes/DataSet_FindAndReplace.py @@ -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", +} diff --git a/classes/DataSet_PathSelector.py b/classes/DataSet_PathSelector.py new file mode 100644 index 0000000..72141f0 --- /dev/null +++ b/classes/DataSet_PathSelector.py @@ -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", +} \ No newline at end of file diff --git a/classes/DataSet_TextFilesLoad.py b/classes/DataSet_TextFilesLoad.py new file mode 100644 index 0000000..406e1cc --- /dev/null +++ b/classes/DataSet_TextFilesLoad.py @@ -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", +} diff --git a/classes/DataSet_TextFilesLoadFromList.py b/classes/DataSet_TextFilesLoadFromList.py new file mode 100644 index 0000000..21e1c3e --- /dev/null +++ b/classes/DataSet_TextFilesLoadFromList.py @@ -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", +} diff --git a/classes/DATASET_TXTFileSaverBatch.py b/classes/DataSet_TextFilesSave.py similarity index 73% rename from classes/DATASET_TXTFileSaverBatch.py rename to classes/DataSet_TextFilesSave.py index ac93717..a29b70e 100644 --- a/classes/DATASET_TXTFileSaverBatch.py +++ b/classes/DataSet_TextFilesSave.py @@ -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", } diff --git a/classes/DATASET_CaptionVisualizer.py b/classes/DataSet_Visualizer.py similarity index 53% rename from classes/DATASET_CaptionVisualizer.py rename to classes/DataSet_Visualizer.py index a2967aa..ded5a57 100644 --- a/classes/DATASET_CaptionVisualizer.py +++ b/classes/DataSet_Visualizer.py @@ -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", } diff --git a/web/util.js b/web/util.js index d460746..da092ee 100644 --- a/web/util.js +++ b/web/util.js @@ -9,4 +9,4 @@ let DataSet = { }, } -app.registerExtension(DataSet); \ No newline at end of file +// app.registerExtension(DataSet); \ No newline at end of file