Create Danbooru Prompt Selector/prompt selector.py

Danbooru Prompt Selector/prompt selector.py

@@ -0,0 +1,194 @@
+import tkinter as tk
+import tkinter.ttk as ttk
+from tkinter import filedialog
+import os
+import csv
+import random
+
+def filter_csv(input_file, output_file, _search_strings):
+    output_directory = os.getcwd()
+    output_file = os.path.join(output_directory, output_file)
+    search_strings = [s.strip() for s in _search_strings.split(',')]
+    with open(input_file, 'r', newline='', encoding='utf-8') as f_in, \
+         open(output_file, 'w', newline='', encoding='utf-8') as f_out:
+        reader = csv.reader(f_in)
+        writer = csv.writer(f_out)
+        writer_count = 0
+        for row in reader:
+            if all(search_str in value for search_str in search_strings for value in row):
+                writer.writerow(row)
+                writer_count += 1
+        return writer_count
+
+def open_file():
+    initial_dir = os.getcwd()
+    filepath = filedialog.askopenfilename(
+        initialdir=initial_dir,
+        filetypes=[("CSV Files", "*.csv")]
+    )
+    if filepath:
+        entry_file_path.delete(0, tk.END)
+        entry_file_path.insert(0, filepath)
+
+def search():
+    global total_rows
+    input_file = entry_file_path.get()
+    keywords = entry_keyword.get()
+    keyword_label.config(text="검색할 키워드: ")
+    keyword_label_mode.config(text="(현재 검색 모드로 동작하고 있습니다.)", fg="blue")
+    output_file = 'txt2img_temp_prompt.csv'
+    writer_count = filter_csv(input_file, output_file, keywords)
+    total_rows = writer_count
+    total_rows_count_label.config(text=f"Total Rows: {writer_count}")
+    text_output.insert(tk.END, f"총 {writer_count}개의 문자열이 검색되어 저장되었습니다.\n")
+    cached_rows = None
+    
+def exclude():
+    global total_rows
+    input_file = entry_file_path.get()
+    keywords = entry_keyword.get()
+    output_file = 'txt2img_temp_prompt.csv'
+    keyword_label.config(text="검색할 키워드: ")
+    keyword_label_mode.config(text="(현재 제외 모드로 동작하고 있습니다.)", fg="red")
+    output_directory = os.getcwd()
+    output_file = os.path.join(output_directory, output_file)
+    search_strings = [s.strip() for s in keywords.split(',')]
+    
+    with open(input_file, 'r', newline='', encoding='utf-8') as f_in, \
+         open(output_file, 'w', newline='', encoding='utf-8') as f_out:
+        reader = csv.reader(f_in)
+        writer = csv.writer(f_out)
+        writer_count = 0
+        for row in reader:
+            if not all(search_str in value for search_str in search_strings for value in row):
+                writer.writerow(row)
+                writer_count += 1
+    total_rows = writer_count
+    total_rows_count_label.config(text=f"Total Rows: {writer_count}")
+    text_output.insert(tk.END, f"총 {writer_count}개의 문자열이 검색되어 저장되었습니다.\n")
+    cached_rows = None
+
+
+def reset():
+    cached_rows = None
+    entry_file_path.delete(0, tk.END)
+    entry_keyword.delete(0, tk.END)
+    entry_deep_search.delete(0, tk.END)
+    text_output.delete('1.0', tk.END)
+
+def random_function():
+    global last_deep_search_keywords, cached_rows
+    current_deep_search_keywords = entry_deep_search.get().strip()
+    if current_deep_search_keywords != last_deep_search_keywords or not cached_rows:
+        with open('txt2img_temp_prompt.csv', 'r', newline='', encoding='utf-8') as f:
+            reader = csv.reader(f)
+            if current_deep_search_keywords:
+                keywords = current_deep_search_keywords.split(',')
+                exclude_keywords = [s[1:].strip() for s in keywords if s.strip().startswith('~')]
+                include_keywords = [s.strip() for s in keywords if not s.strip().startswith('~')]
+                rows = [row for row in reader if not any(exclude in cell for exclude in exclude_keywords for cell in row)]
+                cached_rows = [row for row in rows if all(include in cell for include in include_keywords for cell in row)]
+            else:
+                cached_rows = list(reader)
+        last_deep_search_keywords = current_deep_search_keywords
+    text_output.delete('1.0', tk.END)
+    if cached_rows:
+        random_index = random.randint(0, len(cached_rows) - 1)
+        random_row = cached_rows.pop(random_index)
+        text_output.insert(tk.END, f"{', '.join(random_row)}\n")
+    else:
+        text_output.insert(tk.END, "검색 조건에 맞는 데이터가 없거나 CSV 파일에 데이터가 없습니다.\n")
+
+    cached_rows_count_label.config(text=f"Cached Rows: {len(cached_rows)}")
+
+
+
+
+def copy_to_clipboard():
+    window.clipboard_clear()
+    window.clipboard_append(text_output.get("1.0", tk.END))
+
+def exit_program():
+    window.destroy()
+
+
+def save_settings():
+    with open('app_settings.txt', 'w', encoding='utf-8') as f:
+        f.write(entry_file_path.get() + '\n')
+        f.write(entry_keyword.get() + '\n')
+        f.write(entry_deep_search.get() + '\n')
+
+def load_settings():
+    if os.path.exists('app_settings.txt'):
+        with open('app_settings.txt', 'r', encoding='utf-8') as f:
+            settings = f.readlines()
+            entry_file_path.insert(0, settings[0].strip())
+            entry_keyword.insert(0, settings[1].strip())
+            entry_deep_search.insert(0, settings[2].strip())
+
+def exit_program():
+    save_settings()
+    window.destroy()
+
+window = tk.Tk()
+window.title("Prompt Selector for Danbooru tags")
+last_deep_search_keywords = None
+cached_rows = []
+total_rows = 0
+
+
+# 파일 경로 입력창
+label_file_path = tk.Label(window, text="CSV 파일 경로:")
+label_file_path.grid(row=0, column=0, columnspan=2, sticky='w')
+entry_file_path = tk.Entry(window, width=70)
+entry_file_path.grid(row=1, column=0, columnspan=2, padx=5, pady=5)
+button_open_file = tk.Button(window, text="파일 열기", command=open_file)
+button_open_file.grid(row=1, column=2, padx=5, pady=5)
+
+# 키워드 입력창
+keyword_label = tk.Label(window, text="검색할 키워드: ")
+keyword_label.grid(row=2, column=0, sticky='w')
+entry_keyword = tk.Entry(window, width=70)
+entry_keyword.grid(row=3, column=0, columnspan=2, padx=5, pady=5)
+keyword_label_mode = tk.Label(window, text="")
+keyword_label_mode.grid(row=2, column=1, sticky='w')
+
+# 버튼 프레임
+frame_buttons = tk.Frame(window)
+frame_buttons.grid(row=4, column=0, columnspan=3, pady=10)
+
+# 버튼들
+button_search = tk.Button(frame_buttons, text="검색", command=search)
+button_search.pack(side=tk.LEFT, padx=5)
+button_exclude = tk.Button(frame_buttons, text="제외", command=exclude)
+button_exclude.pack(side=tk.LEFT, padx=5)
+button_reset = tk.Button(frame_buttons, text="리셋", command=reset)
+button_reset.pack(side=tk.LEFT, padx=5)
+
+# 심층검색 키워드 입력창
+label_deep_search = tk.Label(window, text="심층검색 키워드: 키워드 앞에 ~를 붙이면 제외합니다.")
+label_deep_search.grid(row=5, column=0, columnspan=2, sticky='w')
+entry_deep_search = tk.Entry(window, width=70)
+entry_deep_search.grid(row=6, column=0, columnspan=2, padx=5, pady=5)
+
+# Tkinter UI 설정 부분에 레이블 추가
+cached_rows_count_label = tk.Label(window, text="Cached Rows: 0")
+cached_rows_count_label.grid(row=8, column=1, padx=5, pady=5)
+total_rows_count_label = tk.Label(window, text="Total Rows: 0")
+total_rows_count_label.grid(row=8, column=0, padx=5, pady=5)
+
+# 추가 버튼들
+button_random = tk.Button(window, text="랜덤", command=random_function)
+button_random.grid(row=7, column=0, padx=5, pady=5, sticky='w')
+button_copy = tk.Button(window, text="복사", command=copy_to_clipboard)
+button_copy.grid(row=7, column=1, padx=5, pady=5)
+button_exit = tk.Button(window, text="종료", command=exit_program)
+button_exit.grid(row=7, column=2, padx=5, pady=5, sticky='w')
+
+# 출력 텍스트 창
+text_output = tk.Text(window, height=10, width=70)
+text_output.grid(row=9, column=0, columnspan=3, padx=5, pady=5)
+
+load_settings()
+
+window.mainloop()