Update image_to_tag.py

image_to_tag.py

@@ -7,6 +7,7 @@ import threading
 import subprocess
 import sys
 import json
+import re
 
 # Global variables to control the process and track errors
 stop_processing = False
@@ -43,10 +44,11 @@ total_pages = 1  # Initialize total pages
 def update_and_save_config():
     """Update and save the configuration to JSON."""
     save_config_to_json(
-        model="swinv2",  # Default model
+        model="eva02",  # Default model
         general_threshold=general_threshold_var.get(),
         character_threshold=character_threshold_var.get(),
-        model_dir="D:/test/models/wd-swinv2-tagger-v3"  # Default model directory
+        model_dir="D:/test/models/wd-eva02-large-tagger-v3",  # Default model directory
+        caption_mode=caption_mode_var.get()
     )
 
 def show_errors(root):
@@ -77,13 +79,14 @@ def show_errors(root):
 
     error_window.protocol("WM_DELETE_WINDOW", on_close_error_window)
 
-def save_config_to_json(model, general_threshold, character_threshold, model_dir, filepath='config.json'):
+def save_config_to_json(model, general_threshold, character_threshold, model_dir, caption_mode, filepath='config.json'):
     """Save the model and threshold values to a JSON file."""
     config = {
         'model': model,
         'general_threshold': general_threshold,
         'character_threshold': character_threshold,
-        'model_dir': model_dir
+        'model_dir': model_dir,
+        'caption_mode': caption_mode
     }
     try:
         with open(filepath, 'w') as f:
@@ -95,6 +98,7 @@ def open_image_to_tag():
     global stop_processing, error_messages, selected_files, save_directory, caption_window, caption_frame, thumbnails, caption_text_widgets, tag_dict, selected_tag, edit_buttons, tag_text_frame, current_page, total_pages, content_canvas
     global status_var, num_files_var, errors_var, progress, character_threshold_var, general_threshold_var, thread_count_var, batch_size_var
     global start_button, stop_button, prepend_text_var, append_text_var
+    global caption_mode_var  # Khai báo biến toàn cục
 
     # Create Tkinter window
     root = tk.Tk()
@@ -107,16 +111,17 @@ def open_image_to_tag():
     progress = tk.IntVar()
     character_threshold_var = tk.DoubleVar(value=0.35)
     general_threshold_var = tk.DoubleVar(value=0.35)
-    thread_count_var = tk.IntVar(value=4)
-    batch_size_var = tk.IntVar(value=4)
+    thread_count_var = tk.IntVar(value=1)
+    batch_size_var = tk.IntVar(value=8)
     prepend_text_var = tk.StringVar()
     append_text_var = tk.StringVar()
+    caption_mode_var = tk.IntVar(value=1)
     q = queue.Queue()
 
     def center_window(window, width_extra=0, height_extra=0):
         window.update_idletasks()
         width = 100 + width_extra
-        height = 820 + height_extra
+        height = 950 + height_extra
         x = (window.winfo_screenwidth() // 2) - (width // 2)
         y = (window.winfo_screenheight() // 2) - (height // 2)
         window.geometry(f'{width}x{height}+{x}+{y}')
@@ -158,6 +163,7 @@ def open_image_to_tag():
         # Stop button should always be enabled
         stop_button.config(state=tk.NORMAL)
 
+
     def generate_caption(image_path, save_directory, q):
         """Generate captions for a single image using the wd-swinv2-tagger-v3 model."""
         if stop_processing:
@@ -165,13 +171,18 @@ def open_image_to_tag():
 
         try:
             filename = os.path.splitext(os.path.basename(image_path))[0]
-            output_path = os.path.join(save_directory, f"{filename}.txt")
+            output_path = os.path.join(save_directory, f"{filename}.txt")  # Sửa lại tên tệp caption
+
+            # Kiểm tra chế độ tạo caption
+            if caption_mode_var.get() == 2 and os.path.exists(output_path):
+                q.put(image_path)
+                return
 
             command = [
                 sys.executable, 'D:/test/wdv3-timm-main/wdv3_timm.py',
-                '--model', "swinv2",
+                '--model', "eva02",
                 '--image_path', image_path,
-                '--model_dir', "D:/test/models/wd-swinv2-tagger-v3",
+                '--model_dir', "D:/test/models/wd-eva02-large-tagger-v3",
                 '--general_threshold', str(general_threshold_var.get()),
                 '--character_threshold', str(character_threshold_var.get())
             ]
@@ -183,31 +194,54 @@ def open_image_to_tag():
             print(output)  # In ra đầu ra từ lệnh subprocess
             print(error_output)  # In ra đầu ra lỗi từ lệnh subprocess
 
-            # Filter out information to contain only Caption or General tags
-            filtered_output = []
-            recording = False
+            # Lọc thông tin "General tags"
+            general_tags = []
+            recording_general = False
             for line in output.split('\n'):
                 if "General tags" in line:
-                    recording = True
+                    recording_general = True
                     continue
-                if recording:
+                if recording_general:
                     if line.startswith('  '):
                         tag = line.strip().split(':')[0].replace('_', ' ')
-                        filtered_output.append(tag)
+                        general_tags.append(tag)
                     else:
-                        recording = False
-                        break
+                        recording_general = False
+
+            # Lọc thông tin "Character tags"
+            character_tags = []
+            recording_character = False
+            for line in output.split('\n'):
+                if "Character tags" in line:
+                    recording_character = True
+                    continue
+                if recording_character:
+                    if line.startswith('  '):
+                        tag = line.strip().split(':')[0].replace('_', ' ')
 
-            # Convert list of tags to comma-separated string
-            final_tags = ','.join(filtered_output) if filtered_output else "No tags found"
+                        # Loại bỏ từ khóa có chứa từ 'costume'
+                        if 'costume' not in tag.lower():
+                            character_tags.append(tag)  # Giữ lại từ khóa không chứa 'costume'
+                    else:
+                        recording_character = False
+
+            # Kết hợp cả general và character tags
+            final_tags = ','.join(general_tags + character_tags) if general_tags or character_tags else "No tags found"
             print("Filtered output:", final_tags)  # Debug: In ra các nhãn cuối cùng sau khi lọc
 
-            # Add prepend and append text
+            # Thêm văn bản trước và sau
             final_tags = f"{prepend_text_var.get()},{final_tags},{append_text_var.get()}".strip(',')
 
-            # Save result to text file
-            with open(output_path, 'w', encoding='utf-8') as file:
-                file.write(final_tags)
+            # Xử lý ghi đè, nối thêm hoặc bỏ qua caption hiện có
+            if caption_mode_var.get() == 0:  # Overwrite
+                with open(output_path, 'w', encoding='utf-8') as file:
+                    file.write(final_tags)
+            elif caption_mode_var.get() == 1 and os.path.exists(output_path):  # Append
+                with open(output_path, 'a', encoding='utf-8') as file:
+                    file.write(f",{final_tags}")
+            else:  # Tạo mới hoặc ghi đè nếu file không tồn tại
+                with open(output_path, 'w', encoding='utf-8') as file:
+                    file.write(final_tags)
 
             q.put(image_path)
         except Exception as e:
@@ -432,7 +466,8 @@ def open_image_to_tag():
                 file_label.grid(row=i*2, column=1, padx=5, pady=5, sticky="nsew")
 
                 # Check and display caption if available
-                caption_file = os.path.join(save_directory, f"{os.path.basename(file_path)}_tags.txt")
+                filename = os.path.splitext(os.path.basename(file_path))[0]  # Lấy tên tệp không có phần mở rộng
+                caption_file = os.path.join(save_directory, f"{filename}.txt")
                 if os.path.exists(caption_file):
                     with open(caption_file, 'r', encoding='utf-8') as file:
                         caption_text = file.read()
@@ -442,7 +477,7 @@ def open_image_to_tag():
                 caption_text_widget = tk.Text(caption_frame, width=50, height=3, wrap=tk.WORD, font=('Helvetica', 12))
                 caption_text_widget.insert(tk.END, caption_text)
                 caption_text_widget.grid(row=i*2, column=2, padx=5, pady=5, sticky="nsew")
-                caption_text_widget.bind("<FocusOut>", lambda e, fp=file_path: save_caption(fp, caption_text_widget.get("1.0", "end-1c")))
+                caption_text_widget.bind("<FocusOut>", lambda e, fp=file_path, w=caption_text_widget: save_caption(fp, w.get("1.0", "end-1c")))
                 caption_text_widgets.append(caption_text_widget)
 
                 # Update tags in tag_dict
@@ -494,7 +529,8 @@ def open_image_to_tag():
 
     def save_caption(file_path, caption_text):
         """Save caption when user changes it."""
-        output_path = os.path.join(save_directory, f"{os.path.basename(file_path)}_tags.txt")
+        filename = os.path.splitext(os.path.basename(file_path))[0]
+        output_path = os.path.join(save_directory, f"{filename}.txt")
         with open(output_path, 'w', encoding='utf-8') as file:
             file.write(caption_text)
 
@@ -580,7 +616,8 @@ def open_image_to_tag():
 
                 # Update the captions in the respective files
                 for file_path in selected_files:
-                    caption_file = os.path.join(save_directory, f"{os.path.basename(file_path)}_tags.txt")
+                    filename = os.path.splitext(os.path.basename(file_path))[0]
+                    caption_file = os.path.join(save_directory, f"{filename}.txt")
                     if os.path.exists(caption_file):
                         with open(caption_file, 'r', encoding='utf-8') as file:
                             caption_text = file.read()
@@ -614,7 +651,8 @@ def open_image_to_tag():
 
                 # Update the captions in the respective files
                 for i, file_path in enumerate(selected_files):
-                    caption_file = os.path.join(save_directory, f"{os.path.basename(file_path)}_tags.txt")
+                    filename = os.path.splitext(os.path.basename(file_path))[0]
+                    caption_file = os.path.join(save_directory, f"{filename}.txt")
                     if os.path.exists(caption_file):
                         with open(caption_file, 'r', encoding='utf-8') as file:
                             caption_text = file.read()
@@ -644,11 +682,14 @@ def open_image_to_tag():
 
                 # Update the captions in the respective files
                 for i, file_path in enumerate(selected_files):
-                    caption_file = os.path.join(save_directory, f"{os.path.basename(file_path)}_tags.txt")
+                    filename = os.path.splitext(os.path.basename(file_path))[0]
+                    caption_file = os.path.join(save_directory, f"{filename}.txt")
                     if os.path.exists(caption_file):
                         with open(caption_file, 'r', encoding='utf-8') as file:
                             caption_text = file.read()
-                        new_caption_text = caption_text.replace(tag_to_delete, "")
+                        # Remove the tag from the caption
+                        tags = [tag.strip() for tag in caption_text.split(',') if tag.strip() != tag_to_delete]
+                        new_caption_text = ','.join(tags)
                         with open(caption_file, 'w', encoding='utf-8') as file:
                             file.write(new_caption_text)
 
@@ -664,7 +705,8 @@ def open_image_to_tag():
             # Delete the files containing the tag
             files_to_delete = []
             for i, file_path in enumerate(selected_files):
-                caption_file = os.path.join(save_directory, f"{os.path.basename(file_path)}_tags.txt")
+                filename = os.path.splitext(os.path.basename(file_path))[0]
+                caption_file = os.path.join(save_directory, f"{filename}.txt")
                 if os.path.exists(caption_file):
                     with open(caption_file, 'r', encoding='utf-8') as file:
                         caption_text = file.read()
@@ -740,6 +782,19 @@ def open_image_to_tag():
     append_text_entry = tk.Entry(root, textvariable=append_text_var, justify='center', width=20)
     append_text_entry.pack(pady=5)
 
+    # Add Radio buttons for caption mode
+    caption_mode_label = tk.Label(root, text="Caption Mode:")
+    caption_mode_label.pack(fill='x', pady=5)
+
+    overwrite_radio = tk.Radiobutton(root, text="Overwrite existing caption", variable=caption_mode_var, value=0)
+    overwrite_radio.pack(fill='x', pady=5)
+
+    append_radio = tk.Radiobutton(root, text="Append to existing caption", variable=caption_mode_var, value=1)
+    append_radio.pack(fill='x', pady=5)
+
+    skip_radio = tk.Radiobutton(root, text="Skip images with existing caption", variable=caption_mode_var, value=2)
+    skip_radio.pack(fill='x', pady=5)
+
     thread_count_label = tk.Label(root, text="Thread Count:")
     thread_count_label.pack(pady=5)
     thread_count_entry = tk.Entry(root, textvariable=thread_count_var, justify='center', width=5, validate='key')
@@ -775,6 +830,7 @@ def open_image_to_tag():
     general_threshold_var.trace_add('write', lambda *args: update_and_save_config())
     character_threshold_var.trace_add('write', lambda *args: update_and_save_config())
     thread_count_var.trace_add('write', lambda *args: update_and_save_config())
+    caption_mode_var.trace_add('write', lambda *args: update_and_save_config())
 
     center_window(root, width_extra=200)
     root.protocol("WM_DELETE_WINDOW", on_closing)