Danbooru Prompt Selector/prompt selector.py

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 any(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()
    combined_text = entry_fixed_prompt.get() + "\n" + text_output.get("1.0", tk.END)
    window.clipboard_append(combined_text)

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()

def on_ctrl_enter(event):
    random_function()

def on_shift_enter(event):
    random_function()
    copy_to_clipboard()

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)

fixed_prompt_label = tk.Label(window, text="고정 프롬프트: 복사 버튼을 눌렀을 때 동작합니다.")
fixed_prompt_label.grid(row=10, column=0, columnspan=2, sticky='w')
entry_fixed_prompt = tk.Entry(window, width=70)
entry_fixed_prompt.grid(row=11, column=0, columnspan=2, padx=5, pady=5)

shortcut_info_label = tk.Label(window, text="* Ctrl+Enter를 누르면 [랜덤] 버튼이 동작합니다.\n* Shift+Enter를 누르면 [랜덤] 버튼이 동작하고, 그 결과가 클립보드에 복사됩니다.", justify=tk.LEFT)
shortcut_info_label.grid(row=13, column=0, columnspan=3, padx=5, pady=5, sticky='w')

window.bind('<Control-Return>', on_ctrl_enter)
window.bind('<Shift-Return>', on_shift_enter)

load_settings()

window.mainloop()