spider_gui.py

import torch
import torch.nn as nn
from matplotlib import ticker
from torchvision import transforms, models
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from PIL import Image, ImageTk
import os
import random
import shutil
from pathlib import Path
import matplotlib as mpl
from PIL import Image, ImageTk, ImageDraw, ImageFont  # 确保添加了 ImageFont
from datetime import datetime
import time
import glob
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg


import matplotlib.font_manager as fm # 确保导入了字体管理器

# --- Matplotlib 中文乱码处理 ---
# 尝试查找并使用支持中文的字体，增加兼容性
font_list = ['SimHei', 'Microsoft YaHei', 'KaiTi', 'FangSong']
chosen_font = None

for font_name in font_list:
    try:
        # 尝试查找字体，如果找到则停止
        if fm.fontManager.findfont(font_name, fallback_to_default=False):
            chosen_font = font_name
            break
    except:
        continue # 忽略查找失败

# 如果找到了中文字体，则进行设置
if chosen_font:
    # 设置 sans-serif 字体系列的首选字体为找到的中文体
    plt.rcParams['font.sans-serif'] = [chosen_font] + font_list
    # 调试语句：查看实际使用的字体
    print(f"DEBUG: Matplotlib 已配置使用字体: {chosen_font}")
else:
    # 如果所有尝试都失败，至少保留默认字体
    plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
    print("WARNING: 未找到常用中文字体，中文可能显示为方框。")

# 解决保存图像时负号 '-' 显示为方块的问题
plt.rcParams['axes.unicode_minus'] = False


class MultiModelSpiderRecognition:
    def __init__(self, root):
        """
        初始化多模型蜘蛛识别系统。

        Args:
            root (tk.Tk): Tkinter根窗口对象。
        """
        self.root = root
        self.root.title("多模型蜘蛛识别系统")
        # 初始窗口大小
        window_width = 1200
        window_height = 900

        # 获取屏幕的宽度和高度
        screen_width = self.root.winfo_screenwidth()
        screen_height = self.root.winfo_screenheight()

        # 计算窗口的起始位置
        center_x = int(screen_width / 2 - window_width / 2)
        center_y = int(screen_height / 2 - window_height / 2)

        # 设置窗口大小和位置 (居中)
        self.root.geometry(f"{window_width}x{window_height}+{center_x}+{center_y}")
        # 调试语句
        # print(f"DEBUG: 窗口大小: {window_width}x{window_height}, 居中位置: +{center_x}+{center_y}")

        # 使用新的日志系统记录调试信息
        # 注意：log_message 在 setup_gui 之前调用会打印到控制台
        # self.log_message(f"窗口大小: {window_width}x{window_height}, 居中位置: +{center_x}+{center_y}", "DEBUG")

        self.best_model_files = {}  # 记录每个模型的最佳文件
        self.available_models = {}  # 记录所有可用的模型文件
        # 新增：记录已加载模型的 Listbox 显示文本，用于高亮显示
        self.loaded_model_display_texts = set()

        self.models = {}
        self.class_names = []
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.current_model = None

        # 支持的模型列表
        self.supported_models = {
            'ResNet-50': 'resnet50',
            'EfficientNet-B0': 'efficientnet_b0',
            'DenseNet-121': 'densenet121'
        }

        # 为识别结果创建字典，用于存储每个模型的结果文本框
        self.result_texts = {}

        # 图像缩放和平移状态
        self.zoom_level = 1.0  # 缩放级别，1.0 表示“基础缩放”
        self.min_zoom = 0.005  # 最小缩放限制，允许图片缩得更小
        self.max_zoom = 10.0  # 最大缩放限制
        self.pan_x = 0  # 平移 X 坐标
        self.pan_y = 0  # 平移 Y 坐标

        self._last_x = 0  # 鼠标按下时的 X 坐标
        self._last_y = 0  # 鼠标按下时的 Y 坐标
        self.display_width = 400  # 图像预览区域固定宽度
        self.display_height = 400  # 图像预览区域固定高度

        # 新增：用于平滑缩放动画的状态
        self.target_zoom = 1.0  # 缩放的目标级别
        self.is_animating = False  # 标记当前是否正在进行缩放动画

        # 新增：日志记录 Text 组件引用
        self.log_text_widget = None

        self.setup_gui()

    def setup_gui(self):
        """设置GUI界面，优化布局：图片预览固定且更大，图表区域更小，识别结果按模型分隔"""
        # 主框架
        main_frame = ttk.Frame(self.root, padding="5")
        # 将 main_frame 向上扩展，但不能填满窗口底部，因为底部要留给日志区域
        main_frame.pack(fill=tk.BOTH, expand=True, side=tk.TOP)

        # ----------------------------------------------------------------------------------
        # 日志输出区域 (位于窗口最底部，替换了原有的单行状态栏功能)
        log_frame = ttk.LabelFrame(self.root, text="系统日志", padding="0")
        # 将日志区域放在最底部，填充宽度
        log_frame.pack(side=tk.BOTTOM, fill=tk.X, padx=0, pady=0)

        # 日志文本框和滚动条
        scrollbar = ttk.Scrollbar(log_frame)
        # 设置 Text 组件，高度固定为 5 行，使用 Consolas 字体
        self.log_text_widget = tk.Text(log_frame, wrap="word", height=5,
                                       yscrollcommand=scrollbar.set, state='disabled',
                                       font=('Consolas', 8), bg='#f4f4f4')

        scrollbar.config(command=self.log_text_widget.yview)

        # 布局 Text 和 Scrollbar
        scrollbar.pack(side="right", fill="y")
        self.log_text_widget.pack(side="left", fill=tk.BOTH, expand=True)

        # 新增：状态栏 (位于日志区域上方，占据底部的固定高度)
        status_bar = ttk.Frame(self.root, relief=tk.SUNKEN)
        status_bar.pack(side=tk.BOTTOM, fill=tk.X, padx=0, pady=0)

        self.status_var = tk.StringVar(value=">>>准备就绪") # 移动到这里，并作为实例变量

        # 状态栏 Label
        status_label = ttk.Label(status_bar, textvariable=self.status_var,
                                 anchor=tk.W, font=('Arial', 10), background='#dcdcdc')
        status_label.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5, pady=2)

        # ----------------------------------------------------------------------------------
        # 标题
        title_label = ttk.Label(main_frame, text="🕷️ 多模型蜘蛛识别系统",
                                font=("Arial", 20, "bold"))
        title_label.pack(pady=10)

        # 左侧控制面板 (保持不变)
        left_frame = ttk.Frame(main_frame)
        left_frame.pack(side=tk.LEFT, fill=tk.Y, padx=5)

        # 数据准备区域
        data_frame = ttk.LabelFrame(left_frame, text="1. 数据准备", padding="10")
        data_frame.pack(fill=tk.X, pady=5)

        ttk.Label(data_frame, text="数据文件夹:").pack(anchor=tk.W)
        path_frame = ttk.Frame(data_frame)
        path_frame.pack(fill=tk.X, pady=5)

        self.data_path = tk.StringVar()
        ttk.Entry(path_frame, textvariable=self.data_path, width=30).pack(side=tk.LEFT, padx=5)
        ttk.Button(path_frame, text="浏览", command=self.browse_folder).pack(side=tk.LEFT, padx=5)

        ttk.Button(data_frame, text="自动划分数据集",
                   command=self.split_dataset).pack(fill=tk.X, pady=5)

        # 模型训练区域
        train_frame = ttk.LabelFrame(left_frame, text="2. 模型训练", padding="10")
        train_frame.pack(fill=tk.X, pady=5)

        # 模型选择
        ttk.Label(train_frame, text="选择要训练的模型:").pack(anchor=tk.W)
        self.model_vars = {}
        for model_name in self.supported_models.keys():
            var = tk.BooleanVar(value=True)
            self.model_vars[model_name] = var
            ttk.Checkbutton(train_frame, text=model_name, variable=var).pack(anchor=tk.W)

        # 训练参数
        param_frame = ttk.Frame(train_frame)
        param_frame.pack(fill=tk.X, pady=5)

        ttk.Label(param_frame, text="训练轮数:").grid(row=0, column=0, sticky=tk.W)
        self.epochs_var = tk.StringVar(value="20")
        ttk.Entry(param_frame, textvariable=self.epochs_var, width=10).grid(row=0, column=1, padx=5)

        ttk.Label(param_frame, text="批次大小:").grid(row=0, column=2, sticky=tk.W, padx=10)
        self.batch_size_var = tk.StringVar(value="16")
        ttk.Entry(param_frame, textvariable=self.batch_size_var, width=10).grid(row=0, column=3)

        ttk.Button(train_frame, text="开始训练选中模型",
                   command=self.start_training).pack(fill=tk.X, pady=5)

        self.progress_var = tk.DoubleVar()
        self.progress = ttk.Progressbar(train_frame, mode='determinate',
                                        maximum=100, variable=self.progress_var)
        self.progress.pack(fill=tk.X, pady=5)

        # 模型管理区域
        model_frame = ttk.LabelFrame(left_frame, text="3. 模型管理", padding="10")
        model_frame.pack(fill=tk.X, pady=5)

        ttk.Label(model_frame, text="加载已训练模型:").pack(anchor=tk.W)

        # 创建一个框架来容纳Listbox和Scrollbar
        list_scroll_frame = ttk.Frame(model_frame)
        list_scroll_frame.pack(fill=tk.X, pady=5)

        # 创建Listbox - 移除 selectmode=tk.SINGLE，因为我们将用双击来控制状态和颜色
        self.loaded_models_list = tk.Listbox(list_scroll_frame, height=4)
        self.loaded_models_list.pack(side=tk.LEFT, fill=tk.X, expand=True)

        # 绑定双击事件：实现选中/取消高亮和加载/取消加载模型
        self.loaded_models_list.bind('<Double-1>', self.toggle_model_load_status)

        # 创建垂直滚动条
        scrollbar = ttk.Scrollbar(list_scroll_frame, orient=tk.VERTICAL,
                                  command=self.loaded_models_list.yview)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)

        # 将Listbox和Scrollbar连接起来
        self.loaded_models_list.config(yscrollcommand=scrollbar.set)

        # Listbox 不支持 tag_configure，故移除，将在其他函数中直接使用 itemconfigure 设置颜色

        load_frame = ttk.Frame(model_frame)
        load_frame.pack(fill=tk.X)

        ttk.Button(load_frame, text="加载/卸载模型",
                   command=lambda: self.toggle_model_load_status(None, is_button=True)).pack(side=tk.LEFT, padx=2)
        # 删除了 "刷新列表" 按钮
        ttk.Button(load_frame, text="删除模型", command=self.delete_selected_models).pack(side=tk.LEFT, padx=2)

        # 图片识别区域
        predict_frame = ttk.LabelFrame(left_frame, text="4. 图片识别", padding="10")
        predict_frame.pack(fill=tk.X, pady=5)

        ttk.Button(predict_frame, text="选择图片",
                   command=self.select_image).pack(fill=tk.X, pady=2)
        ttk.Button(predict_frame, text="开始识别",
                   command=self.predict_image).pack(fill=tk.X, pady=2)

        # ----------------------------------------------------------------------------------
        # 右侧结果显示区域
        right_frame = ttk.Frame(main_frame)
        right_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True, padx=5)

        # 右侧上部：图片预览 (固定大小) 和 训练图表 (较小)
        top_right_frame = ttk.Frame(right_frame)
        top_right_frame.pack(fill=tk.X, pady=5, expand=False)

        image_width, image_height = 400, 400  # <-- 宽度修改为 600
        image_frame = ttk.LabelFrame(top_right_frame, text="图片预览", padding="10")

        # 强制设置 frame 的像素尺寸，并禁止其尺寸被子组件内容改变
        # 尺寸稍微增加，以预留 LabelFrame 的边框和内部 padding 的空间
        image_frame.config(width=image_width + 20, height=image_height + 40)
        image_frame.pack_propagate(False)  # 阻止子组件（Label）改变父组件（Frame）的大小

        image_frame.pack(side=tk.LEFT, padx=5, fill=tk.BOTH, expand=False)

        # 强制设置Label的字符宽度，以预留空间
        self.image_label = ttk.Label(image_frame, text="请选择图片",
                                     background="white", anchor=tk.CENTER)

        # 绑定鼠标事件用于拖动和平移
        self.image_label.bind("<Button-1>", self._start_pan)  # 鼠标左键按下：记录起始位置
        self.image_label.bind("<B1-Motion>", self._pan_image)  # 鼠标左键拖动：平移

        # 绑定鼠标滚轮事件用于缩放 (Windows/Linux 是 <MouseWheel>，macOS 是 <Button-4>/<Button-5>)
        # Tkinter 在 Windows/Linux 上统一将滚轮识别为 <MouseWheel>
        self.image_label.bind("<MouseWheel>", self._zoom_image)
        # 针对部分 Linux 系统或特定配置，也绑定 Button-4/Button-5
        self.image_label.bind("<Button-4>", lambda event: self._zoom_image(event, delta=1))  # 向上
        self.image_label.bind("<Button-5>", lambda event: self._zoom_image(event, delta=-1))  # 向下

        # 将 Label 的 pack 属性改为不 expand，因为它现在由父 Frame 固定大小
        # 同时增加 padx/pady 确保其在 LabelFrame 内部居中或填充
        self.image_label.pack(fill=tk.BOTH, expand=False, pady=5, padx=5)

        # 训练历史图表 - 占用剩余空间
        chart_frame = ttk.LabelFrame(top_right_frame, text="训练历史图表", padding="5")
        chart_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=5)

        # 设置图表大小 (宽高比约为 1.2:1)
        self.fig, self.ax = plt.subplots(2, 1, figsize=(6, 4), dpi=100)
        self.canvas = FigureCanvasTkAgg(self.fig, chart_frame)
        self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True, pady=5, padx=5)
        # 初始化图表
        self.update_training_chart({})
        self.canvas.draw()

        # 右侧下部：识别结果 (按模型划分)
        result_frame = ttk.LabelFrame(right_frame, text="识别结果", padding="10")
        result_frame.pack(fill=tk.BOTH, expand=True, pady=5)

        # 使用三个独立的Frame来显示每个模型的结果
        model_result_frame = ttk.Frame(result_frame)
        model_result_frame.pack(fill=tk.BOTH, expand=True)

        # 为每个支持的模型创建一个子区域
        for model_name in self.supported_models.keys():
            # 创建一个子框架
            frame = ttk.LabelFrame(model_result_frame, text=model_name, padding="5")
            frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=5, pady=5)

            # 创建文本框显示结果
            text_widget = tk.Text(frame, wrap=tk.WORD, height=10, width=1)  # width=1 允许它在pack时扩展
            scrollbar = ttk.Scrollbar(frame, orient=tk.VERTICAL, command=text_widget.yview)

            text_widget.configure(yscrollcommand=scrollbar.set)

            text_widget.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
            scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
            self.result_texts[model_name] = text_widget

        self.refresh_model_list()

    def browse_folder(self):
        """选择数据文件夹"""
        folder = filedialog.askdirectory(title="选择包含蜘蛛类别文件夹的目录")
        if folder:
            self.data_path.set(folder)
            self.scan_classes(folder)

    def scan_classes(self, folder_path):
        """扫描类别"""
        classes = []
        for item in Path(folder_path).iterdir():
            if item.is_dir():
                classes.append(item.name)

        if classes:
            self.status_var.set(f"发现 {len(classes)} 个类别: {', '.join(classes)}")
            self.class_names = classes
        else:
            self.status_var.set("未找到类别文件夹")

    def validate_data_structure(self):
        """验证数据目录结构是否正确"""
        data_path = Path(self.data_path.get())

        # 检查是否有不应该的文件夹名
        invalid_folders = ['train_val', 'train', 'val', 'test']
        found_classes = []

        for item in data_path.iterdir():
            if item.is_dir():
                if item.name in invalid_folders:
                    messagebox.showerror(
                        "数据目录错误",
                        f"发现无效文件夹: {item.name}\n\n"
                        f"请确保{data_path}目录下直接是蜘蛛类别文件夹\n"
                        f"例如: 黑寡妇/, 跳蛛/, 狼蛛/\n"
                        f"而不是: train_val/, train/, val/"
                    )
                    return False
                found_classes.append(item.name)

        if len(found_classes) < 2:
            messagebox.showerror(
                "数据不足",
                f"只找到 {len(found_classes)} 个类别: {found_classes}\n"
                f"至少需要2个类别才能训练模型"
            )
            return False

        # 调试语句已修改为使用 self.log_message
        self.log_message(f"数据目录正确！找到类别: {found_classes}", "DEBUG")
        return True

    def split_dataset(self):
        """
        自动划分数据集：将选定文件夹下的图片按 8:2 比例划分为 train 和 val 文件夹。

        Returns:
            bool: 划分成功返回 True，失败返回 False。
        """
        data_path = self.data_path.get()
        if not data_path:
            self.root.after(0, lambda: messagebox.showerror("错误", "请先选择数据文件夹！"))
            return False

        base_path = Path(data_path)
        # 目标路径
        target_dir = base_path / 'train_val'

        # 划分比例
        TRAIN_RATIO = 0.8

        try:
            self.status_var.set(">>>开始划分数据集...")
            self.log_message("开始划分数据集...", "INFO")

            # 清理旧的划分结果
            if target_dir.exists():
                shutil.rmtree(target_dir)

            # 创建新的 train 和 val 目录
            train_dir = target_dir / 'train'
            val_dir = target_dir / 'val'
            train_dir.mkdir(parents=True, exist_ok=True)
            val_dir.mkdir(parents=True, exist_ok=True)

            # 获取所有子文件夹（类别）
            all_classes = [d for d in base_path.iterdir() if d.is_dir() and d.name != 'train_val']

            # 调试语句
            # self.log_message(f"发现类别文件夹: {[c.name for c in all_classes]}", "DEBUG")

            if not all_classes:
                self.root.after(0,
                                lambda: messagebox.showerror("错误", f"在 {base_path} 中未找到任何类别子文件夹！"))
                self.status_var.set(">>>划分失败")
                return False

            total_files_moved = 0

            for class_dir in all_classes:
                class_name = class_dir.name

                # 确保目标类文件夹存在
                (train_dir / class_name).mkdir(exist_ok=True)
                (val_dir / class_name).mkdir(exist_ok=True)

                # 获取所有图片文件
                all_files = list(class_dir.glob('*.jpg')) + list(class_dir.glob('*.png'))
                random.shuffle(all_files)  # 随机打乱

                if not all_files:
                    # print() 替换为 self.log_message()
                    self.log_message(f"类别 '{class_name}' 下未找到图片，跳过。", "WARNING")
                    continue

                # 计算划分数量
                train_count = int(len(all_files) * TRAIN_RATIO)

                train_files = all_files[:train_count]
                val_files = all_files[train_count:]

                # 移动文件
                for f in train_files:
                    shutil.copy(f, train_dir / class_name / f.name)
                for f in val_files:
                    shutil.copy(f, val_dir / class_name / f.name)

                total_files_moved += len(all_files)

            if total_files_moved == 0:
                self.root.after(0, lambda: messagebox.showerror("错误",
                                                                "未移动任何文件，请确保所选文件夹及其子文件夹中有JPG/PNG图片！"))
                self.status_var.set(">>>划分失败")
                return False

            self.root.after(0, lambda: self.status_var.set(f">>>数据集划分完成！共处理 {total_files_moved} 个文件。"))
            self.log_message(f"数据集划分完成！共处理 {total_files_moved} 个文件。", "SUCCESS")
            return True  # 划分成功

        except Exception as e:
            self.root.after(0, lambda: messagebox.showerror("划分失败", f"数据集划分过程中发生错误: {str(e)}"))
            self.status_var.set(">>>划分失败")
            self.log_message(f"数据集划分过程中发生错误: {str(e)}", "ERROR")
            return False  # 划分失败

    def log_message(self, message: str, level: str = "INFO"):
        """
        统一的日志记录方法，将日志写入 Text 组件并自动滚动到底部。

        Args:
            message (str): 要记录的日志信息。
            level (str): 日志级别 (如 INFO, ERROR, WARNING, SUCCESS, DEBUG)。
        """
        # 兼容性：检查日志组件是否已创建 (setup_gui 之前会是 None)
        if self.log_text_widget is None:
            print(f"[{datetime.now().strftime('%H:%M:%S')}] [{level}] {message}")
            return

        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        log_entry = f"[{timestamp}] [{level}] {message}\n"

        # 1. 解锁并插入日志
        self.log_text_widget.config(state='normal')
        self.log_text_widget.insert(tk.END, log_entry)

        # 2. 自动滚动到最新日志（底部）
        self.log_text_widget.see(tk.END)

        # 3. 重新锁定 Text 组件，防止用户编辑
        self.log_text_widget.config(state='disabled')

        # 4. 同时更新兼容性的状态栏
        self.status_var.set(f"[{level}] {message}")

    def create_model(self, model_name, num_classes):
        """创建指定模型"""
        if model_name == 'ResNet-50':
            model = models.resnet50(pretrained=True)
            model.fc = nn.Linear(model.fc.in_features, num_classes)
        elif model_name == 'EfficientNet-B0':
            model = models.efficientnet_b0(pretrained=True)
            model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
        elif model_name == 'DenseNet-121':
            model = models.densenet121(pretrained=True)
            model.classifier = nn.Linear(model.classifier.in_features, num_classes)
        else:
            raise ValueError(f"不支持的模型: {model_name}")

        return model.to(self.device)

    def update_training_chart(self, new_data):
        """
        更新训练过程图表 (损失和准确率)。

        Args:
            new_data (dict): 包含模型名和训练历史数据的字典。
                             例如: {'model_name_timestamp': {'train_losses': [...], 'val_accuracies': [...]}}
        """
        # 确保图表对象已初始化
        if self.ax is None:
            self.log_message("图表对象未初始化，跳过更新。", "WARNING")
            return

        # 清除旧图
        self.ax[0].clear()  # 损失图
        self.ax[1].clear()  # 准确率图

        # ------------------- 绘制损失图 -------------------
        self.ax[0].set_title('训练损失 (Loss)', fontsize=10)
        self.ax[0].set_xlabel('Epoch', fontsize=8)
        self.ax[0].set_ylabel('Loss', fontsize=8)
        self.ax[0].tick_params(axis='both', which='major', labelsize=7)
        self.ax[0].grid(True, linestyle='--', alpha=0.6)

        # ------------------- 绘制准确率图 -------------------
        self.ax[1].set_title('验证准确率 (Accuracy)', fontsize=10)
        self.ax[1].set_xlabel('Epoch', fontsize=8)
        self.ax[1].set_ylabel('Accuracy (%)', fontsize=8)
        self.ax[1].tick_params(axis='both', which='major', labelsize=7)
        self.ax[1].grid(True, linestyle='--', alpha=0.6)

        # 遍历所有模型数据并绘制
        for model_label, data in new_data.items():
            epochs = range(1, len(data['train_losses']) + 1)

            # **修改点 1：为损失曲线添加 label**
            self.ax[0].plot(epochs, data['train_losses'], label=f'Loss ({model_label})',
                            linestyle='-', marker='.', markersize=3)

            # **修改点 2：为准确率曲线添加 label**
            self.ax[1].plot(epochs, data['val_accuracies'], label=f'Acc ({model_label})',
                            linestyle='-', marker='.', markersize=3)

        # **修改点 3：移动并调用 legend()**
        # 只有在绘制完带 label 的曲线后，调用 legend() 才能显示内容。

        # 损失图例
        # 确保只在有曲线时调用 legend，如果图表清空后立即调用 legend 可能会有空警告。
        # 这里使用 has_data 隐式检查，因为 new_data 是非空的
        if new_data:
            self.ax[0].legend(fontsize=8, loc='upper right')  # 第 599 行
            self.ax[1].legend(fontsize=8, loc='lower right')  # 第 605 行

        # 优化 x 轴刻度，避免刻度过多
        for ax in self.ax:
            if ax.get_lines():
                max_epoch = max(len(ax.get_lines()[0].get_xdata()), 1)
                if max_epoch > 10:
                    ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True, nbins=10))
                else:
                    ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True))

        # 刷新画布
        self.canvas.draw_idle()

    def train_single_model(self, model_name, data_path=None):
        """
        训练单个模型。

        该函数执行完整的模型训练流程，包括数据加载、模型初始化、
        训练循环、验证、模型保存，并实时更新 GUI 上的训练损失、
        准确率图表和进度条。

        Args:
            model_name (str): 要训练的模型名称。
            data_path (str, optional): 训练数据的路径。默认为 None，将从 self.data_path 获取。

        Returns:
            tuple: (best_acc, train_losses, val_accuracies, timestamp)
        """
        from torch.utils.data import DataLoader
        from torchvision import datasets

        # 生成时间戳
        timestamp = time.strftime("%Y%m%d_%H%M%S")

        # 路径设置
        if data_path is None:
            data_path = Path(self.data_path.get()) / "train_val"
        else:
            data_path = Path(data_path)

        # 数据预处理
        train_transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.RandomHorizontalFlip(0.3),
            transforms.RandomRotation(10),
            transforms.ColorJitter(brightness=0.2, contrast=0.2),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

        val_transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

        # 加载数据
        train_path = data_path / 'train'
        val_path = data_path / 'val'

        try:
            train_dataset = datasets.ImageFolder(root=train_path, transform=train_transform)
            val_dataset = datasets.ImageFolder(root=val_path, transform=val_transform)
        except Exception as e:
            raise RuntimeError(f"数据加载失败，请检查路径和图片文件: {str(e)}")

        # **修复点 1: 在这里获取准确的类别名称和数量**
        # 确保模型创建和数据加载使用的类别数量一致 (解决 Target out of bounds 错误)

        # 必须先检查数据集是否为空
        if len(train_dataset) == 0:
            raise RuntimeError(f"训练集为空！请检查 {train_path} 路径下是否有图片。")
        if len(val_dataset) == 0:
            raise RuntimeError(f"验证集为空！请检查 {val_path} 路径下是否有图片。")

        # 获取准确的类别信息
        class_names_local = train_dataset.classes
        num_classes_local = len(class_names_local)

        # **重要：将准确的类别名称更新到实例变量，供预测时使用**
        self.class_names = class_names_local

        # 调试语句：检查加载数据量
        # print() 替换为 self.log_message()
        self.log_message(
            f"训练集大小: {len(train_dataset)}, 验证集大小: {len(val_dataset)}, 类别数: {num_classes_local}", "DEBUG")

        batch_size = int(self.batch_size_var.get())
        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

        # 创建模型
        # **修复点 2: 使用本地准确的 num_classes_local 创建模型**
        model = self.create_model(model_name, num_classes_local)

        # 训练设置
        criterion = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

        epochs = int(self.epochs_var.get())
        best_acc = 0
        train_losses = []
        val_accuracies = []

        for epoch in range(epochs):
            # 训练阶段
            model.train()
            epoch_train_loss = 0

            if len(train_loader) == 0:
                # print() 替换为 self.log_message()
                self.log_message("train_loader 为空，跳过训练阶段。", "WARNING")
                break

            for batch_idx, (inputs, targets) in enumerate(train_loader):
                inputs, targets = inputs.to(self.device), targets.to(self.device)

                optimizer.zero_grad()
                outputs = model(inputs)
                loss = criterion(outputs, targets)
                loss.backward()
                optimizer.step()

                epoch_train_loss += loss.item()

                # --- 进度条更新代码 (保持不变) ---
                current_batch_progress = (batch_idx + 1) / len(train_loader)
                overall_progress_percent = ((epoch) + current_batch_progress) * 100 / epochs

                self.root.after(0, lambda p=overall_progress_percent: self.progress_var.set(p))
                # --- 进度条更新代码结束 ---

            # 计算平均训练损失
            if len(train_loader) > 0:
                avg_train_loss = epoch_train_loss / len(train_loader)
            else:
                avg_train_loss = 0.0
            train_losses.append(avg_train_loss)

            # 验证阶段
            model.eval()
            correct = 0
            total = 0

            # 必须检查验证集是否为空，防止除以零
            if len(val_loader) == 0:
                acc = 0.0
            else:
                with torch.no_grad():
                    for inputs, targets in val_loader:
                        inputs, targets = inputs.to(self.device), targets.to(self.device)
                        outputs = model(inputs)
                        _, predicted = outputs.max(1)
                        total += targets.size(0)
                        correct += predicted.eq(targets).sum().item()
                acc = 100. * correct / total

            val_accuracies.append(acc)

            # 实时更新图表
            current_data = {f"{model_name}_{timestamp}": {
                'train_losses': train_losses,
                'val_accuracies': val_accuracies
            }}
            self.root.after(0, lambda: self.update_training_chart(current_data))

            self.root.after(0, lambda: self.status_var.set(
                f"{model_name} - Epoch {epoch + 1}/{epochs} - 损失: {avg_train_loss:.4f} - 准确率: {acc:.1f}%"
            ))

            if acc > best_acc:
                best_acc = acc
                # 保存模型 - 使用时间戳命名
                model_filename = f'{model_name}_model_{timestamp}.pth'
                model_info = {
                    'model_state_dict': model.state_dict(),
                    'class_names': class_names_local,  # **修复点 3: 保存准确的类别名称**
                    'accuracy': acc,
                    'epoch': epoch,
                    'train_losses': train_losses.copy(),
                    'val_accuracies': val_accuracies.copy(),
                    'timestamp': timestamp,
                    'model_name': model_name,
                    'total_epochs': epochs
                }
                torch.save(model_info, model_filename)
                self.log_message(f"模型 {model_filename} 保存成功，准确率: {acc:.1f}%", "INFO")

                # 更新最佳模型记录
                self.best_model_files[model_name] = model_filename

            scheduler.step()

        # 训练完成后，确保进度条显示 100%
        self.root.after(0, lambda: self.progress_var.set(100))

        return best_acc, train_losses, val_accuracies, timestamp

    # 完整的子函数 (精简回核心功能)
    def start_training(self):
        """
        启动选中模型的训练过程。将训练过程放入单独线程，避免阻塞 GUI。
        """
        data_path = self.data_path.get()
        if not data_path:
            messagebox.showerror("错误", "请先选择数据文件夹！")
            return

        # 检查是否有模型被选中
        selected_models = [name for name, var in self.model_vars.items() if var.get()]
        if not selected_models:
            messagebox.showerror("错误", "请至少选择一个要训练的模型！")
            return

        self.status_var.set(">>>正在准备训练环境...")

        import threading

        # 训练逻辑函数
        def training_thread_target():
            # 1. 确保数据集划分是最新且有效的
            # **修改点：直接检查 split_dataset 的返回值**
            if not self.split_dataset():
                # split_dataset 失败时会自行弹窗提示
                self.root.after(0, lambda: self.status_var.set(">>>训练因数据集划分失败而停止。"))
                return

            try:
                # 检查划分结果的文件夹路径
                train_val_path = Path(data_path) / "train_val"
                # 再次检查 train 文件夹是否确实存在，这能捕捉到 split_dataset 内部的失败
                if not (train_val_path / 'train').exists():
                    self.root.after(0, lambda: messagebox.showerror("划分错误", "数据集划分失败，未找到训练目录。"))
                    return

                # 清空训练历史记录
                self.all_training_data = {}
                self.root.after(0, lambda: self.update_training_chart({}))
                self.root.after(0, lambda: self.progress_var.set(0))

                for model_name in selected_models:
                    self.root.after(0, lambda m=model_name: self.status_var.set(f">>>开始训练模型: {m}..."))

                    # 调用实际训练函数
                    best_acc, _, _, _ = self.train_single_model(model_name, data_path=train_val_path)

                    self.root.after(0, lambda m=model_name, acc=best_acc: self.status_var.set(
                        f">>>模型 {m} 训练完成，最佳准确率: {acc:.1f}%"
                    ))

            except RuntimeError as e:
                self.root.after(0, lambda err=str(e): messagebox.showerror("训练失败", f"训练失败: {err}"))
                self.root.after(0, lambda: self.status_var.set(">>>训练因错误停止。"))
            except Exception as e:
                self.root.after(0, lambda err=str(e): messagebox.showerror("未知错误", f"程序遇到未知错误: {err}"))
                self.root.after(0, lambda: self.status_var.set(">>>训练因未知错误停止。"))
            finally:
                # 所有模型训练完毕或发生错误后
                self.root.after(0, lambda: self.refresh_model_list())
                self.root.after(0, lambda: self.status_var.set(">>>训练任务结束。"))

        # 创建并启动线程
        training_thread = threading.Thread(target=training_thread_target)
        training_thread.start()

    def refresh_model_list(self):
        """刷新模型列表 - 增强版本，并保持已加载模型的颜色。

        该函数会扫描当前目录下所有模型文件，并将其信息展示在Listbox中。
        如果模型已被加载（在self.loaded_model_display_texts中），则高亮显示。
        """
        self.loaded_models_list.delete(0, tk.END)
        self.available_models.clear()

        # 查找所有模型文件
        model_files = glob.glob('*_model_*.pth')
        self.log_message(f"发现 {len(model_files)} 个模型文件。", "DEBUG")

        for model_file in model_files:
            try:
                # 从文件名解析信息
                if 'ResNet-50_model_' in model_file:
                    model_name = 'ResNet-50'
                elif 'EfficientNet-B0_model_' in model_file:
                    model_name = 'EfficientNet-B0'
                elif 'DenseNet-121_model_' in model_file:
                    model_name = 'DenseNet-121'
                else:
                    continue

                # 加载模型信息获取准确率
                checkpoint = torch.load(model_file, map_location='cpu')
                accuracy = checkpoint.get('accuracy', 0)
                timestamp = checkpoint.get('timestamp', '未知时间')

                display_text = f"{model_name} | 准确率: {accuracy:.1f}% | 时间: {timestamp}"

                # 插入 Listbox
                self.loaded_models_list.insert(tk.END, display_text)
                self.available_models[display_text] = model_file

                # 检查是否是已加载的模型，并应用颜色
                if display_text in self.loaded_model_display_texts:
                    # 获取刚刚插入的项的索引
                    index = self.loaded_models_list.size() - 1
                    # 直接设置背景和前景颜色
                    self.loaded_models_list.itemconfigure(index, {'background': '#3498db', 'foreground': 'white'})

                self.log_message(f"加载模型信息成功: {display_text}", "DEBUG")

            except Exception as e:
                # print() 替换为 self.log_message()
                self.log_message(f"加载模型文件 {model_file} 失败: {e}", "ERROR")

    def clear_models(self):
        """清除所有加载的模型。

        该函数会清空加载的模型字典，重置状态变量，并清除Listbox中的所有高亮显示。
        """
        # 清除 Listbox 中的高亮
        for i in range(self.loaded_models_list.size()):
            # 恢复默认背景和前景
            self.loaded_models_list.itemconfigure(i, {'background': '', 'foreground': ''})

        self.models = {}
        self.current_model = None
        self.loaded_model_display_texts.clear()  # 清除记录
        self.status_var.set("已清除所有加载的模型")

    def delete_selected_models(self):
        """
        删除Listbox中选中的模型文件，并从磁盘上永久移除。
        删除前会弹出确认窗口。
        """
        # 获取选中的索引
        selection = self.loaded_models_list.curselection()
        if not selection:
            messagebox.showwarning("警告", "请先在列表中选择一个要删除的模型！")
            return

        # 选中的是 Listbox 中的索引，通常只有一个，我们只处理第一个
        index = selection[0]
        display_text = self.loaded_models_list.get(index)
        model_file_name = self.available_models.get(display_text)
        model_name = self._get_model_name_from_display_text(display_text)

        if not model_file_name:
            self.log_message(f"模型文件路径丢失: {display_text}", "ERROR")
            messagebox.showerror("错误", "无法找到对应的模型文件路径。")
            return

        # 弹出确认对话框
        confirm = messagebox.askyesno(
            "确认删除",
            f"您确定要永久删除模型文件吗？\n\n模型: {model_name}\n文件: {model_file_name}\n\n此操作不可撤销！"
        )

        if confirm:
            try:
                # 1. 从磁盘删除文件
                if os.path.exists(model_file_name):
                    os.remove(model_file_name)
                    self.log_message(f"已永久删除模型文件: {model_file_name}", "SUCCESS")
                else:
                    self.log_message(f"文件不存在: {model_file_name}，但已从列表中移除。", "WARNING")

                # 2. 从内部记录和 Listbox 中移除
                # 尝试卸载（如果已加载）
                if display_text in self.loaded_model_display_texts:
                    self.loaded_model_display_texts.remove(display_text)
                    self.models.pop(model_name, None)
                    # 检查并清除 current_model 引用
                    if self.current_model == model_name:
                        self.current_model = None

                # 从 Listbox 中删除项
                self.loaded_models_list.delete(index)
                # 从 available_models 中删除记录
                self.available_models.pop(display_text, None)

                self.status_var.set(f"已成功删除模型: {model_name}")
                self.log_message(f"模型 {model_name} 已从系统和磁盘中移除。", "INFO")

            except Exception as e:
                messagebox.showerror("删除失败", f"删除文件 {model_file_name} 时发生错误: {str(e)}")
                self.log_message(f"删除模型失败: {str(e)}", "ERROR")

    def select_image(self):
        """选择图片，并重置视图状态"""
        self.current_image_path = filedialog.askopenfilename(
            title="选择蜘蛛图片",
            filetypes=[("图片文件", "*.jpg *.jpeg *.png *.bmp")]
        )

        if self.current_image_path:
            try:
                # 关键：将原始 PIL 图像保存在实例变量中
                self.original_pil_image = Image.open(self.current_image_path).convert('RGB')

                # 重置视图状态
                self.zoom_level = 1.0
                self.pan_x = 0
                self.pan_y = 0

                # 调用统一的图像更新函数
                self._update_displayed_image(self.original_pil_image)

                self.status_var.set(f"已选择图片: {os.path.basename(self.current_image_path)}")

            except Exception as e:
                # 错误修正：将 selfbox.showerror 替换为 messagebox.showerror
                messagebox.showerror("错误", f"加载图片失败: {str(e)}")

    def _start_pan(self, event):
        """
        处理鼠标左键按下事件，记录平移起始点。
        """
        self._last_x = event.x
        self._last_y = event.y
        # 鼠标按下时，将焦点设置给 Label，确保滚轮事件能在 Label 上触发
        self.image_label.focus_set()

    def _pan_image(self, event):
        """
        处理鼠标左键拖动事件，计算平移偏移量并更新图片显示。
        """
        if self.original_pil_image is None:
            return

        # 计算新的偏移量
        dx = event.x - self._last_x
        dy = event.y - self._last_y

        self.pan_x += dx
        self.pan_y += dy

        # 限制平移范围，防止图片完全移出视野 (可选，但推荐)

        # 获取当前缩放后图片的实际尺寸
        w, h = self.original_pil_image.size
        scaled_w = int(w * self.zoom_level)
        scaled_h = int(h * self.zoom_level)

        # 计算最大平移边界 (这里简化处理，允许拖动直到边缘)
        max_pan_x = max(0, scaled_w - self.display_width) / 2
        max_pan_y = max(0, scaled_h - self.display_height) / 2

        # 限制平移范围
        self.pan_x = max(-max_pan_x, min(self.pan_x, max_pan_x))
        self.pan_y = max(-max_pan_y, min(self.pan_y, max_pan_y))

        self._last_x = event.x
        self._last_y = event.y

        self._update_displayed_image(self.original_pil_image)

    def _zoom_image(self, event, delta=None):
        """
        处理鼠标滚轮事件，实现缩放，并在缩放后重置平移状态，以保持视图稳定。

        Args:
            event (tk.Event): Tkinter事件对象。
            delta (int, optional): 对于 Button-4/Button-5 事件，手动指定缩放方向 (1: 放大, -1: 缩小)。
        """
        if self.original_pil_image is None:
            return

        # 设置缩放步长：使用 1.05 替代 1.1，实现更丝滑的缩放效果
        SMOOTH_FACTOR = 1.05

        # 统一处理滚轮方向：Windows/Linux 的 <MouseWheel> 事件
        if event.num == 0 or event.delta != 0:
            # event.delta > 0: 向上滚动 (放大)
            # event.delta < 0: 向下滚动 (缩小)
            factor = SMOOTH_FACTOR if event.delta > 0 else (1 / SMOOTH_FACTOR)
        elif delta is not None:
            # macOS/Linux 的 Button-4/Button-5 事件
            factor = SMOOTH_FACTOR if delta > 0 else (1 / SMOOTH_FACTOR)
        else:
            return

        new_zoom = self.zoom_level * factor

        # 限制缩放级别：使用 self.min_zoom 和 self.max_zoom
        # 确保缩放后不会超出定义的范围
        new_zoom = max(self.min_zoom, min(new_zoom, self.max_zoom))

        if new_zoom != self.zoom_level:
            self.zoom_level = new_zoom

            # **关键修正：删除或注释掉以下两行，不再重置平移参数**
            # self.pan_x = 0
            # self.pan_y = 0

            # 调试语句已替换为 self.log_message()
            # 这里的 DEBUG 消息也需要修改，因为我们不再重置 Pan
            self.log_message(f"Wheel event. New zoom level set to {self.zoom_level:.4f}. Pan retained.", "DEBUG")

            # 调用统一的图像更新函数
            self._update_displayed_image(self.original_pil_image)

    def _update_displayed_image(self, pil_image: Image.Image):
        """
        根据当前缩放和平移状态，处理 PIL 图像并在 self.image_label 上显示。

        Args:
            pil_image (Image.Image): 要处理的 PIL 图像对象。
        """
        if pil_image is None:
            return

        # 1. 计算保持比例的初始缩放 (使原图最佳适应 self.display_width x self.display_height)
        original_width, original_height = pil_image.size
        width_ratio = self.display_width / original_width
        height_ratio = self.display_height / original_height

        # 基础缩放比例 (确保图片最初完整显示)
        base_scale = min(width_ratio, height_ratio)

        # 最终应用于 PIL 图像的缩放比例
        final_scale = base_scale * self.zoom_level

        # 缩放后的总尺寸
        scaled_w = int(original_width * final_scale)
        scaled_h = int(original_height * final_scale)

        # 2. 缩放图像
        scaled_w = max(1, scaled_w)
        scaled_h = max(1, scaled_h)
        scaled_image = pil_image.resize((scaled_w, scaled_h), Image.Resampling.LANCZOS)

        # 3. 平移限制：如果图片小于显示区域，则不能平移
        # 只有当 scaled_w > display_width 时，才允许平移 X 轴
        pan_limit_x = max(0, (scaled_w - self.display_width) / 2)
        pan_limit_y = max(0, (scaled_h - self.display_height) / 2)

        # 重新限制 pan_x 和 pan_y
        self.pan_x = max(-pan_limit_x, min(self.pan_x, pan_limit_x))
        self.pan_y = max(-pan_limit_y, min(self.pan_y, pan_limit_y))

        # 4. 创建 $W \times H$ 画布并居中/裁剪粘贴
        display_image = Image.new('RGB', (self.display_width, self.display_height), 'white')

        # 裁剪的左上角坐标，相对于 scaled_image 的 (0, 0)
        final_crop_x_min = 0
        final_crop_y_min = 0

        # 粘贴到 display_image 上的左上角坐标
        paste_x_offset = 0
        paste_y_offset = 0

        # --- 核心修正逻辑 ---

        # X 轴处理 (宽度)
        if scaled_w <= self.display_width:
            # 图片宽度小于或等于显示区域：居中
            paste_x_offset = (self.display_width - scaled_w) // 2
            final_crop_x_min = 0  # 裁剪从 0 开始
            # pan_x 在这种情况下应为 0，已在上面被限制
        else:
            # 图片宽度大于显示区域：需要裁剪，根据 pan_x 决定裁剪起点
            final_crop_x_min = scaled_w // 2 - self.display_width // 2 - self.pan_x
            final_crop_x_min = max(0, final_crop_x_min)  # 确保不越界
            paste_x_offset = 0  # 裁剪后的图片贴到 (0, 0)

        # Y 轴处理 (高度)
        if scaled_h <= self.display_height:
            # 图片高度小于或等于显示区域：居中
            paste_y_offset = (self.display_height - scaled_h) // 2
            final_crop_y_min = 0  # 裁剪从 0 开始
            # pan_y 在这种情况下应为 0，已在上面被限制
        else:
            # 图片高度大于显示区域：需要裁剪，根据 pan_y 决定裁剪起点
            final_crop_y_min = scaled_h // 2 - self.display_height // 2 - self.pan_y
            final_crop_y_min = max(0, final_crop_y_min)  # 确保不越界
            paste_y_offset = 0  # 裁剪后的图片贴到 (0, 0)

        # 裁剪框的右下角坐标
        final_crop_x_max = final_crop_x_min + (self.display_width - paste_x_offset)
        final_crop_y_max = final_crop_y_min + (self.display_height - paste_y_offset)

        # 确保裁剪框不会超出 scaled_image 的右下边界
        final_crop_x_max = min(final_crop_x_max, scaled_w)
        final_crop_y_max = min(final_crop_y_max, scaled_h)

        # 5. 执行裁剪/粘贴
        # 裁剪出所需的部分
        cropped_image = scaled_image.crop((final_crop_x_min, final_crop_y_min, final_crop_x_max, final_crop_y_max))

        # 粘贴到 display_image 上，使用计算出的 offset 来实现居中
        display_image.paste(cropped_image, (paste_x_offset, paste_y_offset))

        # 6. 显示到 Label
        photo = ImageTk.PhotoImage(display_image)
        self.image_label.configure(image=photo)
        self.image_label.image = photo
    def predict_image(self):
        """
        对当前选择的图片进行多模型预测。
        识别成功后，在图片预览上绘制边界框作为视觉反馈，并标注所有模型的Top-3结果。
        """
        if not self.current_image_path:
            messagebox.showerror("错误", "请先选择一张图片")
            return

        if not self.models:
            messagebox.showerror("错误", "请先加载至少一个已训练模型")
            return

        if not self.class_names:
            messagebox.showerror("错误", "无法获取类别名称，请确保已加载正确的模型")
            return

        self.status_var.set("正在识别...")

        try:
            # 预测所需的图像预处理
            transform = transforms.Compose([
                transforms.Resize((224, 224)),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])

            image = Image.open(self.current_image_path).convert('RGB')
            input_tensor = transform(image).unsqueeze(0).to(self.device)

            # 存储所有已加载模型的 Top-3 预测结果
            all_top_predictions = {}

            # 为每个加载的模型进行预测
            for model_name, model in self.models.items():
                with torch.no_grad():
                    outputs = model(input_tensor)
                    probabilities = torch.nn.functional.softmax(outputs[0], dim=0)

                # --- 提取前三名结果 ---
                probabilities_list = probabilities.cpu().numpy()
                top_indices = probabilities_list.argsort()[-3:][::-1]  # 获取前三高的索引并反转

                current_top_predictions_list = []
                for i in top_indices:
                    class_name = self.class_names[i]
                    confidence_percent = probabilities[i].item() * 100
                    current_top_predictions_list.append((class_name, confidence_percent))

                # 存储该模型的结果
                all_top_predictions[model_name] = current_top_predictions_list

                # --- 更新文本框结果 ---
                confidence, predicted = torch.max(probabilities, 0)
                predicted_class = self.class_names[predicted.item()]
                confidence_percent = confidence.item() * 100

                # 显示结果
                result_text = f"识别结果: {predicted_class}\n"
                result_text += f"置信度: {confidence_percent:.2f}%\n\n"
                result_text += "所有类别概率:\n"

                # 显示所有类别的概率
                for i, prob in enumerate(probabilities):
                    result_text += f"  {self.class_names[i]}: {prob.item() * 100:.2f}%\n"

                # 更新对应模型的文本框
                text_widget = self.result_texts[model_name]
                text_widget.delete(1.0, tk.END)
                text_widget.insert(1.0, result_text)

            # --- 识别成功后，在图片上绘制框图和标签 ---
            original_image = Image.open(self.current_image_path).convert('RGB')

            # 调用辅助函数来绘制框图和所有模型的标签
            if all_top_predictions:
                self.update_image_with_box(original_image, all_top_predictions)

            self.status_var.set("识别完成！")

        except Exception as e:
            self.status_var.set(f"识别失败: {str(e)}")
            messagebox.showerror("错误", f"识别过程中出错: {str(e)}")

    def update_image_with_box(self, image: Image.Image, all_top_predictions: dict, line_color="red", line_width=5):
        """
        在图片上绘制一个代表识别区域的框图，并在右上角堆叠标注所有模型的识别结果。

        Args:
            image (Image.Image): 原始PIL图片对象。
            all_top_predictions (dict): 字典，键为模型名称，值为包含 (class_name, confidence_percent) 元组的列表。
            line_color (str): 边界框颜色。
            line_width (int): 边界框线宽。
        """
        from PIL import ImageDraw, ImageFont

        # 获取图片尺寸
        img_width, img_height = image.size
        # 确保使用 RGBA 模式以便绘制半透明背景
        if image.mode != 'RGBA':
            image = image.convert('RGBA')
        draw = ImageDraw.Draw(image)

        # 绘制矩形框 (X: 25% - 75%, Y: 25% - 75%)
        x_min = int(img_width * 0.25)
        y_min = int(img_height * 0.25)
        x_max = int(img_width * 0.75)
        y_max = int(img_height * 0.75)

        # 绘制粗线矩形
        for i in range(line_width):
            # draw.rectangle 不支持 line_width，需要多次绘制
            draw.rectangle([x_min - i, y_min - i, x_max + i, y_max + i], outline=line_color)

        # --- 绘制所有模型的排行名称和置信度 ---

        # 尝试加载中文字体，并将字体大小设置为 20
        font_size = 15
        try:
            # 路径可能需要根据系统调整，这里使用一个常见的Windows字体
            font = ImageFont.truetype("C:\\Windows\\Fonts\\simsun.ttc", font_size)
        except IOError:
            try:
                # 尝试一个更通用的英文/数字字体
                font = ImageFont.truetype("arial.ttf", font_size)
            except IOError:
                font = ImageFont.load_default()

        # 标签起始位置：框的右侧
        initial_text_x = x_max + 10  # 框的右侧稍微偏移
        current_text_y = y_min  # 初始 Y 位置与框的顶部对齐

        # 循环遍历所有模型的预测结果
        for model_name, top_predictions in all_top_predictions.items():

            # 构建标签文本
            label_text = f"来自 {model_name} (Top-3):\n"
            for i, (class_name, confidence) in enumerate(top_predictions):
                label_text += f"  {i + 1}. {class_name} ({confidence:.2f}%)\n"

            # 获取文本的边界框 (bbox)
            try:
                # 使用 draw.textbbox
                bbox = draw.textbbox((initial_text_x, current_text_y), label_text, font=font)
            except Exception:
                # 兼容旧PIL/Pillow版本
                try:
                    text_w, text_h = draw.textsize(label_text, font=font)
                    bbox = [initial_text_x, current_text_y, initial_text_x + text_w, current_text_y + text_h]
                except Exception:
                    # 最终回退
                    text_w, text_h = 400, 100  # 估算值
                    bbox = [initial_text_x, current_text_y, initial_text_x + text_w, current_text_y + text_h]

            # 增加一些填充
            padding = 5
            text_bg_rect = [bbox[0] - padding, bbox[1] - padding, bbox[2] + padding, bbox[3] + padding]

            # 绘制黑色半透明背景 (RGBA)
            draw.rectangle(text_bg_rect, fill=(0, 0, 0, 180))

            # 绘制文本 (将颜色更改为亮黄色 (255, 255, 0))
            draw.text((initial_text_x, current_text_y), label_text, font=font, fill=(255, 255, 0))

            # 更新下一个标签的起始 Y 位置：在当前标签的底部下方留出 2 像素间隔 (进一步优化信息密度)
            current_text_y = bbox[3] + 2
        # --- 绘制代码结束 ---

        # 将图片转回 RGB 模式
        image = image.convert('RGB')

        # 关键修改：将带有绘制内容的图片设置为新的 self.original_pil_image
        self.original_pil_image = image

        # 调用统一的图像更新函数，应用当前的缩放和平移状态
        self._update_displayed_image(self.original_pil_image)

    def toggle_model_load_status(self, event=None, is_button=False):
        """
        双击 Listbox 中的模型项时触发：
        如果模型已加载，则卸载模型并取消高亮。
        如果模型未加载，则加载模型并高亮显示（同时卸载同名的旧版本模型）。

        Args:
            event (tk.Event, optional): Tkinter事件对象。
            is_button (bool): 如果是通过“加载模型”按钮触发，则为True。
        """
        # 确定选中的索引
        if event:
            # 双击事件发生时，获取点击处的索引
            index = self.loaded_models_list.nearest(event.y)
        else:
            # 如果是按钮触发，尝试获取当前 Listbox 中光标/焦点的索引
            selection = self.loaded_models_list.curselection()
            if not selection:
                messagebox.showwarning("警告", "请先选择一个模型")
                return
            index = selection[0]

        display_text = self.loaded_models_list.get(index)
        model_file = self.available_models.get(display_text)

        if not model_file:
            messagebox.showerror("错误", "模型文件路径丢失")
            return

        model_name = self._get_model_name_from_display_text(display_text)

        # 核心逻辑：检查该模型是否已加载
        if display_text in self.loaded_model_display_texts:
            # --- 状态：已加载 -> 动作：卸载并取消高亮 ---

            # 从内部记录中移除
            self.models.pop(model_name, None)
            # 确保 current_model 指向已卸载的模型时被清除
            if self.current_model == model_name:
                self.current_model = None

            self.loaded_model_display_texts.remove(display_text)

            # 取消 Listbox 的高亮颜色
            self.loaded_models_list.itemconfigure(index, {'background': '', 'foreground': ''})

            # 更新状态栏
            self.status_var.set(f"已卸载: {model_name} (文件: {os.path.basename(model_file)})")

        else:
            # --- 状态：未加载 -> 动作：加载并高亮显示 ---

            # 1. 清除同名模型的旧高亮和记录（如果存在）
            old_display_text = next((text for text in self.loaded_model_display_texts
                                     if self._get_model_name_from_display_text(text) == model_name), None)

            if old_display_text:
                self.loaded_model_display_texts.remove(old_display_text)
                # 遍历 Listbox 清除旧项的颜色
                for i in range(self.loaded_models_list.size()):
                    item_text = self.loaded_models_list.get(i)
                    if item_text == old_display_text:
                        self.loaded_models_list.itemconfigure(i, {'background': '', 'foreground': ''})
                        break

            # 2. 调用原有的 load_model 核心逻辑
            try:
                checkpoint = torch.load(model_file, map_location=self.device)

                # 创建/加载模型
                num_classes = len(checkpoint['class_names'])
                # 确保类别名称已更新
                if not self.class_names or self.class_names != checkpoint['class_names']:
                    self.class_names = checkpoint['class_names']

                model = self.create_model(model_name, num_classes)
                model.load_state_dict(checkpoint['model_state_dict'])
                model.eval()

                self.models[model_name] = model
                self.current_model = model_name

                accuracy = checkpoint.get('accuracy', 0)

                # 3. 应用新高亮并更新记录
                self.loaded_models_list.itemconfigure(index, {'background': '#3498db', 'foreground': 'white'})
                self.loaded_model_display_texts.add(display_text)

                # 4. 更新状态栏 (不弹窗)
                self.status_var.set(f"已加载: {model_name} (准确率: {accuracy:.1f}%)")

            except Exception as e:
                messagebox.showerror("错误", f"加载模型失败: {str(e)}")
                # 如果加载失败，确保 Listbox 项没有被错误高亮
                self.loaded_models_list.itemconfigure(index, {'background': '', 'foreground': ''})
                if display_text in self.loaded_model_display_texts:
                    self.loaded_model_display_texts.remove(display_text)

    def _get_model_name_from_display_text(self, display_text):
        """
        从Listbox的显示文本中提取模型名称，如 'ResNet-50 | 准确率: 85.0% | 时间: ...' -> 'ResNet-50'
        """
        # 使用 split(' | ')，只取第一部分
        return display_text.split(' | ')[0]

def main():
    root = tk.Tk()
    app = MultiModelSpiderRecognition(root)
    root.mainloop()


if __name__ == "__main__":
    main()