Upload AI 圖庫管理器.py

AI 圖庫管理器.py

@@ -0,0 +1,304 @@
+import os
+import shutil
+import json
+import collections
+import threading
+import tkinter as tk
+from tkinter import filedialog
+from PIL import Image, ImageTk, PngImagePlugin
+import customtkinter as ctk
+from tkinterdnd2 import TkinterDnD, DND_FILES
+
+# 設定主題
+ctk.set_appearance_mode("Dark")
+ctk.set_default_color_theme("blue")
+
+# 防止 PIL 截斷過長的 Metadata
+PngImagePlugin.MAX_TEXT_CHUNK = 5 * 1024 * 1024 
+
+class ImageViewer(ctk.CTkToplevel):
+    """獨立視窗：圖片檢視器"""
+    def __init__(self, parent, image_path, prompt_text):
+        super().__init__(parent)
+        self.title("圖片檢視器 - 僅顯示正向 Prompt")
+        self.geometry("900x750")
+        self.attributes("-topmost", True)
+        
+        self.img_frame = ctk.CTkFrame(self)
+        self.img_frame.pack(fill="both", expand=True, padx=10, pady=10)
+        
+        self.img_label = ctk.CTkLabel(self.img_frame, text="讀取中...")
+        self.img_label.pack(fill="both", expand=True)
+
+        self._load_image(image_path)
+
+        # 底部 Prompt 資訊區
+        self.info_frame = ctk.CTkFrame(self, height=150)
+        self.info_frame.pack(fill="x", padx=10, pady=(0, 10))
+        
+        ctk.CTkLabel(self.info_frame, text="正向 Prompt:", font=("Arial", 12, "bold")).pack(anchor="w", padx=5, pady=5)
+        
+        self.prompt_box = ctk.CTkTextbox(self.info_frame, height=120)
+        self.prompt_box.pack(fill="x", padx=5, pady=5)
+        self.prompt_box.insert("1.0", prompt_text)
+        self.prompt_box.configure(state="disabled")
+
+    def _load_image(self, path):
+        try:
+            pil_img = Image.open(path)
+            base_height = 550
+            w_percent = (base_height / float(pil_img.size[1]))
+            w_size = int((float(pil_img.size[0]) * float(w_percent)))
+            
+            pil_img = pil_img.resize((w_size, base_height), Image.Resampling.LANCZOS)
+            self.ctk_img = ctk.CTkImage(light_image=pil_img, dark_image=pil_img, size=(w_size, base_height))
+            self.img_label.configure(image=self.ctk_img, text="")
+        except Exception as e:
+            self.img_label.configure(text=f"圖片載入失敗: {e}")
+
+class AIImageManagerV5(ctk.CTk, TkinterDnD.DnDWrapper):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        
+        try:
+            self.TkdndVersion = TkinterDnD._require(self)
+        except Exception as e:
+            print(f"DnD Error: {e}")
+
+        self.title("AI 圖庫管理器 v5.0 - 專注正向 Prompt")
+        self.geometry("1280x850")
+        
+        self.image_index = [] 
+        self.prompt_counts = collections.Counter()
+        
+        self.context_menu = tk.Menu(self, tearoff=0)
+        self.context_menu.add_command(label="📥 另存圖片", command=self._download_current_image)
+        self.selected_image_path = None
+
+        self._setup_ui()
+
+    def _setup_ui(self):
+        self.grid_columnconfigure(1, weight=1)
+        self.grid_rowconfigure(0, weight=1)
+
+        # === 左側側邊欄 ===
+        self.sidebar = ctk.CTkFrame(self, width=300, corner_radius=0)
+        self.sidebar.grid(row=0, column=0, sticky="nsew")
+        self.sidebar.grid_rowconfigure(4, weight=1)
+
+        self.drop_label = ctk.CTkLabel(self.sidebar, text="📂 拖放資料夾至此", 
+                                      fg_color="gray25", corner_radius=8, height=80)
+        self.drop_label.grid(row=0, column=0, padx=10, pady=(20, 10), sticky="ew")
+        
+        try:
+            self.drop_label.drop_target_register(DND_FILES)
+            self.drop_label.dnd_bind('<<Drop>>', self._on_folder_drop)
+        except: pass
+
+        self.search_var = ctk.StringVar()
+        self.search_var.trace("w", self._on_search_change)
+        ctk.CTkEntry(self.sidebar, placeholder_text="🔍 搜尋正向關鍵字...", textvariable=self.search_var)\
+            .grid(row=1, column=0, padx=10, pady=10, sticky="ew")
+
+        # 統計設定
+        ctrl_frame = ctk.CTkFrame(self.sidebar, fg_color="transparent")
+        ctrl_frame.grid(row=2, column=0, padx=10, pady=5, sticky="ew")
+        
+        ctk.CTkLabel(ctrl_frame, text="熱門標籤:", font=("Arial", 14, "bold")).pack(side="left")
+        
+        self.top_k_var = ctk.StringVar(value="20")
+        self.top_k_menu = ctk.CTkOptionMenu(ctrl_frame, values=["20", "50", "100"], 
+                                           variable=self.top_k_var, width=80,
+                                           command=self._update_stats_limit)
+        self.top_k_menu.pack(side="right")
+
+        self.stats_scroll = ctk.CTkScrollableFrame(self.sidebar, label_text="")
+        self.stats_scroll.grid(row=4, column=0, padx=10, pady=10, sticky="nsew")
+
+        # === 右側預覽區 ===
+        self.gallery_frame = ctk.CTkScrollableFrame(self, label_text="圖片預覽區")
+        self.gallery_frame.grid(row=0, column=1, padx=20, pady=20, sticky="nsew")
+
+    # --- 核心邏輯修改區 ---
+
+    def _clean_a1111_prompt(self, raw_text):
+        """
+        過濾器：只保留 Negative prompt 之前的內容
+        """
+        if not raw_text: return ""
+        
+        # 關鍵邏輯：一旦出現 Negative prompt，切斷後面所有內容
+        if "Negative prompt:" in raw_text:
+            cleaned = raw_text.split("Negative prompt:")[0]
+        else:
+            cleaned = raw_text
+
+        # 額外清理：有時候 A1111 就算沒有 Negative prompt，也會有 Steps: 等資訊
+        # 如果您希望連 Steps 資訊都不要，可以再加一行：
+        # if "Steps:" in cleaned: cleaned = cleaned.split("Steps:")[0]
+        
+        return cleaned.strip()
+
+    def _extract_metadata(self, img_path):
+        """讀取並清理 Metadata"""
+        raw_data = ""
+        
+        # 1. 嘗試讀取同名 txt
+        txt_path = os.path.splitext(img_path)[0] + ".txt"
+        if os.path.exists(txt_path):
+            try:
+                with open(txt_path, 'r', encoding='utf-8') as tf:
+                    raw_data = tf.read().strip()
+            except: pass
+        
+        # 2. 如果 txt 沒東西，讀 PNG
+        if not raw_data:
+            try:
+                img = Image.open(img_path)
+                info = img.info or {}
+                
+                # A1111
+                if 'parameters' in info:
+                    raw_data = info['parameters']
+                # ComfyUI (還是保留，以免混用)
+                elif 'prompt' in info:
+                    raw_data = info['prompt'] # JSON 格式
+                elif 'workflow' in info:
+                    raw_data = info['workflow']
+            except: pass
+            
+        # 3. 執行清理邏輯 (只留 Positive)
+        return self._clean_a1111_prompt(raw_data)
+
+    def _on_folder_drop(self, event):
+        path = event.data
+        if path.startswith('{') and path.endswith('}'): path = path[1:-1]
+        if os.path.isdir(path):
+            threading.Thread(target=self._process_directory, args=(path,), daemon=True).start()
+
+    def _process_directory(self, folder_path):
+        self.image_index = []
+        self.prompt_counts.clear()
+        
+        valid_exts = ('.png', '.jpg', '.jpeg', '.webp')
+        try:
+            files = [f for f in os.listdir(folder_path) if f.lower().endswith(valid_exts)]
+        except: return
+        
+        total = len(files)
+        
+        for idx, f in enumerate(files):
+            img_path = os.path.join(folder_path, f)
+            
+            # 這裡拿到的已經是「乾淨的正向 Prompt」
+            clean_prompt = self._extract_metadata(img_path)
+            
+            # --- 統計邏輯 ---
+            # 簡單清理一些標點符號，讓統計更準確
+            words = clean_prompt.replace('\n', ',').replace('(', '').replace(')', '').split(',')
+            
+            unique_words = set(w.strip().lower() for w in words if w.strip())
+            
+            # 過濾常見無意義詞彙
+            stop_words = {'masterpiece', 'best quality', 'highres', 'absurdres', '4k', '8k'} 
+            filtered_words = [w for w in unique_words if len(w) > 2 and w not in stop_words]
+            
+            self.prompt_counts.update(filtered_words)
+            
+            self.image_index.append({
+                "path": img_path,
+                "prompt": clean_prompt.lower(), 
+                "display_prompt": clean_prompt, 
+                "mtime": os.path.getmtime(img_path)
+            })
+            
+            if idx % 10 == 0:
+                self.after(0, lambda t=f"索引中... {idx}/{total}": self.title(t))
+
+        self.after(0, self._ui_update_after_index)
+
+    def _ui_update_after_index(self):
+        self.title(f"AI 圖庫管理器 - {len(self.image_index)} 張")
+        self._render_stats_buttons()
+        self._filter_and_render_gallery()
+
+    def _update_stats_limit(self, value):
+        self._render_stats_buttons()
+
+    def _render_stats_buttons(self):
+        for w in self.stats_scroll.winfo_children(): w.destroy()
+        
+        try: limit = int(self.top_k_var.get())
+        except: limit = 20
+            
+        for word, count in self.prompt_counts.most_common(limit):
+            disp = (word[:22] + "...") if len(word) > 25 else word
+            btn = ctk.CTkButton(self.stats_scroll, text=f"{disp} ({count})", 
+                          fg_color="transparent", border_width=1, anchor="w", height=24,
+                          command=lambda w=word: self.search_var.set(w))
+            btn.pack(pady=2, fill="x")
+
+    def _on_search_change(self, *args):
+        self._filter_and_render_gallery(self.search_var.get().lower())
+
+    def _filter_and_render_gallery(self, keyword=""):
+        for w in self.gallery_frame.winfo_children(): w.destroy()
+        
+        if not keyword:
+            filtered = sorted(self.image_index, key=lambda x: x['mtime'], reverse=True)[:20]
+            label_text = "最新生成的 20 張圖片"
+        else:
+            filtered = sorted([i for i in self.image_index if keyword in i['prompt']], 
+                              key=lambda x: x['mtime'], reverse=True)[:50]
+            label_text = f"搜尋 '{keyword}' 結果: {len(filtered)} 筆"
+
+        self.gallery_frame.configure(label_text=label_text)
+
+        if not filtered:
+            ctk.CTkLabel(self.gallery_frame, text="無圖片").pack(pady=20)
+            return
+
+        columns = 4
+        for idx, data in enumerate(filtered):
+            try:
+                pil_img = Image.open(data['path'])
+                pil_img.thumbnail((180, 180))
+                ctk_img = ctk.CTkImage(light_image=pil_img, dark_image=pil_img, size=pil_img.size)
+                
+                card = ctk.CTkFrame(self.gallery_frame, fg_color="gray20")
+                card.grid(row=idx//columns, column=idx%columns, padx=5, pady=5, sticky="nsew")
+                
+                lbl = ctk.CTkLabel(card, image=ctk_img, text="")
+                lbl.pack(padx=5, pady=5)
+                
+                lbl.bind("<Button-1>", lambda e, p=data['path'], txt=data['display_prompt']: self._open_lightbox(p, txt))
+                lbl.bind("<Button-3>", lambda e, p=data['path']: self._show_context_menu(e, p))
+                
+                fname = os.path.basename(data['path'])
+                ctk.CTkLabel(card, text=fname[:15]+"..." if len(fname)>15 else fname, 
+                            font=("Arial", 10)).pack(pady=(0,5))
+            except Exception as e:
+                print(e)
+
+    def _open_lightbox(self, path, prompt_text):
+        ImageViewer(self, path, prompt_text)
+
+    def _show_context_menu(self, event, path):
+        self.selected_image_path = path
+        try:
+            self.context_menu.tk_popup(event.x_root, event.y_root)
+        finally:
+            self.context_menu.grab_release()
+
+    def _download_current_image(self):
+        if not self.selected_image_path: return
+        src_path = self.selected_image_path
+        dest_path = filedialog.asksaveasfilename(
+            defaultextension=".png", initialfile=os.path.basename(src_path),
+            filetypes=[("PNG", "*.png"), ("All", "*.*")])
+        if dest_path:
+            shutil.copy2(src_path, dest_path)
+
+if __name__ == "__main__":
+    app = AIImageManagerV5()
+    app.mainloop()