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, 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, ',') word_counter = Counter() tag_cooccurrences = Counter() for content in contents: try: 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 processing content: {e}") 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) G = nx.Graph() tags_cooccurrence = defaultdict(int) for content in contents: try: 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 processing content: {e}") 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) plt.figure(figsize=(24, 12)) gradio_blue = '#0B0F19' plt.gca().set_facecolor(gradio_blue) degrees = dict(G.degree) node_size = [v * 100 for v in degrees.values()] node_color = [degrees[n] for n in G.nodes] 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_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') 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.close() print("Network graph saved as", output_network_graph_file) 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') 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.xlabel('Tag') plt.ylabel('Frequency') plt.xticks(rotation=45, ha='right') plt.tight_layout() output_table_file = join(output_dir, 'tag_frequency_table.png') plt.savefig(output_table_file) plt.close() print("Tag frequency table saved as", output_table_file) return output_wordcloud_file, output_network_graph_file, output_table_file 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 DataSet_Visualizer: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "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}) }, } INPUT_IS_LIST = True RETURN_TYPES = ("STRING", "IMAGE") RETURN_NAMES = ("GraphsPaths", "GraphsImages") OUTPUT_IS_LIST = (True, True) FUNCTION = "Visualize" OUTPUT_NODE = True CATEGORY = "🔶DATASET🔶" def Visualize(self, TextFileContents, Seperator, WordCloudTop, NetworkGraphTop, FrequencyGraphTop): try: wc, ng, fg = generate_wordcloud_and_network_graph(TextFileContents, Seperator[0], WordCloudTop[0], NetworkGraphTop[0], FrequencyGraphTop[0]) images = [] if os.path.exists(wc): images.append(load_image(wc)) if os.path.exists(ng): images.append(load_image(ng)) if os.path.exists(fg): images.append(load_image(fg)) return ([wc, ng, fg], images,) except Exception as e: print(f"Error saving: {e}") return (["", "", ""],) N_CLASS_MAPPINGS = { "DataSet_Visualizer": DataSet_Visualizer, } N_DISPLAY_NAME_MAPPINGS = { "DataSet_Visualizer": "DataSet_Visualizer", }