Initial Commit
This commit is contained in:
daxcay
@@ -0,0 +1,155 @@
|
||||
import os
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image, ImageSequence
|
||||
import os
|
||||
from os.path import isfile, join
|
||||
import matplotlib.pyplot as plt
|
||||
from wordcloud import WordCloud
|
||||
import networkx as nx
|
||||
from collections import Counter
|
||||
import pandas as pd
|
||||
|
||||
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:
|
||||
with open(file_path, 'r') as f:
|
||||
tags = f.read().strip().split(',')
|
||||
word_counter.update(tags)
|
||||
for tag1 in tags:
|
||||
for tag2 in tags:
|
||||
if tag1 != tag2:
|
||||
tag_cooccurrences[(tag1, tag2)] += 1
|
||||
|
||||
# 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)
|
||||
|
||||
plt.figure(figsize=(16, 9))
|
||||
plt.imshow(wordcloud, interpolation='bilinear')
|
||||
plt.axis('off')
|
||||
|
||||
output_wordcloud_file = join(output_dir, 'wordcloud.png')
|
||||
plt.savefig(output_wordcloud_file, bbox_inches='tight', pad_inches=0)
|
||||
plt.close()
|
||||
|
||||
print("Word cloud saved as", output_wordcloud_file)
|
||||
|
||||
# Network graph generation
|
||||
G = nx.Graph()
|
||||
|
||||
top_tags = [tag for tag, _ in word_counter.most_common(top_n_network)]
|
||||
|
||||
for (tag1, tag2), weight in tag_cooccurrences.items():
|
||||
if tag1 in top_tags and tag2 in top_tags:
|
||||
G.add_edge(tag1, tag2, weight=weight)
|
||||
|
||||
plt.figure(figsize=(16, 9))
|
||||
pos = nx.kamada_kawai_layout(G)
|
||||
|
||||
nx.draw(G, pos, with_labels=True, font_size=10, node_color='skyblue', node_size=2000, edge_color='gray', linewidths=1, font_weight='bold')
|
||||
edge_labels = {(tag1, tag2): weight for (tag1, tag2), weight in tag_cooccurrences.items() if tag1 in top_tags and tag2 in top_tags}
|
||||
nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_color='red')
|
||||
|
||||
plt.margins(0)
|
||||
|
||||
output_network_graph_file = join(output_dir, 'network_graph.png')
|
||||
plt.savefig(output_network_graph_file, bbox_inches='tight', pad_inches=0)
|
||||
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.reset_index(inplace=True)
|
||||
tag_freq_table.columns = ['Tag', 'Frequency']
|
||||
|
||||
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", join(output_dir, 'tag_frequency_table.png'))
|
||||
|
||||
return output_wordcloud_file, output_network_graph_file, join(output_dir, 'tag_frequency_table.png')
|
||||
|
||||
def pilToImage(image):
|
||||
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
||||
|
||||
|
||||
def load_image(image_path):
|
||||
|
||||
try:
|
||||
img = Image.open(image_path)
|
||||
image = img.convert("RGB")
|
||||
loaded_image = pilToImage(image)
|
||||
except Exception as e:
|
||||
print(f"Error loading image from '{image_path}': {e}")
|
||||
|
||||
return loaded_image
|
||||
|
||||
|
||||
def create_empty_image(width=100, height=100, color=(255, 255, 255)):
|
||||
try:
|
||||
empty_image = Image.new("RGB", (width, height), color)
|
||||
except Exception as e:
|
||||
print(f"Error creating empty image: {e}")
|
||||
empty_image = None
|
||||
|
||||
return empty_image
|
||||
|
||||
|
||||
class DRMN_CaptionVisualizer:
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"TextFilePathList": ("STRING", {"forceInput": True}),
|
||||
"WordCloudTop": ("INT", {"default": 1, "min": 1, "max": 9999}),
|
||||
"NetworkGraphTop": ("INT", {"default": 1, "min": 1, "max": 9999}),
|
||||
"FrequencyTableTop": ("INT", {"default": 1, "min": 1, "max": 9999})
|
||||
},
|
||||
}
|
||||
|
||||
INPUT_IS_LIST = True
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("Images",)
|
||||
OUTPUT_IS_LIST = (True,)
|
||||
FUNCTION = "Visualize"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "JDCN Dataset Tools"
|
||||
|
||||
def Visualize(self, TextFilePathList, WordCloudTop, NetworkGraphTop, FrequencyTableTop):
|
||||
|
||||
try:
|
||||
|
||||
directory_path = os.path.dirname(TextFilePathList[0])
|
||||
visualize_path = os.path.join(directory_path, "visualize")
|
||||
os.makedirs(visualize_path, exist_ok=True)
|
||||
wc, ng, fg = generate_wordcloud_and_network_graph(TextFilePathList, visualize_path, WordCloudTop[0], NetworkGraphTop[0], FrequencyTableTop[0])
|
||||
return ([wc, ng, fg],)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error saving: {e}")
|
||||
|
||||
return (["", "", ""],)
|
||||
|
||||
|
||||
N_CLASS_MAPPINGS = {
|
||||
"DRMN_CaptionVisualizer": DRMN_CaptionVisualizer,
|
||||
}
|
||||
|
||||
N_DISPLAY_NAME_MAPPINGS = {
|
||||
"DRMN_CaptionVisualizer": "DRMN_CaptionVisualizer",
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
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 DRMN_TXTFileSaver:
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"content": ("STRING",{"forceInput": True}),
|
||||
"file_name": ("STRING",{"forceInput": True}),
|
||||
"directory": ("STRING", {"default": "directory path"}),
|
||||
"mode": (['Merge','Overwrite','SaveNew','MergeAndSaveNew'],),
|
||||
},
|
||||
}
|
||||
|
||||
# INPUT_IS_LIST = True
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "SaveIT"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
CATEGORY = "RAFTOOLS"
|
||||
|
||||
def SaveIT(self, content, file_name, directory, mode):
|
||||
try:
|
||||
|
||||
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 ()
|
||||
|
||||
|
||||
N_CLASS_MAPPINGS = {
|
||||
"DRMN_TXTFileSaver": DRMN_TXTFileSaver,
|
||||
}
|
||||
|
||||
N_DISPLAY_NAME_MAPPINGS = {
|
||||
"DRMN_TXTFileSaver": "DRMN_TXTFileSaver",
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
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 DRMN_TagManipulatorByImageNames:
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"ImageNames": ("STRING", {"forceInput": True}),
|
||||
"TagsDirectory": ("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 = "RAFTOOLS"
|
||||
|
||||
def SaveIT(self, ImageNames, TagsDirectory, Captions, Backup):
|
||||
try:
|
||||
|
||||
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 ()
|
||||
|
||||
|
||||
N_CLASS_MAPPINGS = {
|
||||
"DRMN_TagManipulatorByImageNames": DRMN_TagManipulatorByImageNames,
|
||||
}
|
||||
|
||||
N_DISPLAY_NAME_MAPPINGS = {
|
||||
"DRMN_TagManipulatorByImageNames": "DRMN_TagManipulatorByImageNames",
|
||||
}
|
||||
Reference in New Issue
Block a user