1530 lines
61 KiB
Python
1530 lines
61 KiB
Python
import os
|
||
import json
|
||
import cv2
|
||
import numpy as np
|
||
import customtkinter as ctk
|
||
from tkinter import filedialog, messagebox
|
||
from PIL import Image
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from typing import List, Optional, Tuple
|
||
from tkinterdnd2 import DND_FILES, TkinterDnD
|
||
import threading
|
||
import time
|
||
import pickle
|
||
|
||
from collections import Counter
|
||
|
||
from net.resnet_embedding import create_resnet50_embedding
|
||
from settings import settings
|
||
|
||
# 设置customtkinter的外观
|
||
ctk.set_appearance_mode("System")
|
||
ctk.set_default_color_theme("blue")
|
||
|
||
|
||
class EmbeddingFoodClassifierApp:
|
||
def __init__(self, root):
|
||
self.root = root
|
||
self.root.title("数字味道-食物识别系统 (Embedding版)")
|
||
self.root.geometry("1400x800")
|
||
|
||
# 设备设置
|
||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
print(f"使用设备: {self.device}")
|
||
|
||
# 当前上传的图片列表
|
||
self.uploaded_images = []
|
||
self.current_results = []
|
||
|
||
# 识别时间记录
|
||
self.recognition_start_time = None
|
||
self.recognition_duration = 0
|
||
|
||
# 模型和FAISS索引相关
|
||
self.model = None
|
||
self.db_vectors = None
|
||
self.db_matrix = None
|
||
self.image_paths = []
|
||
self.labels = []
|
||
self.class_info = {}
|
||
self.class_names = []
|
||
self.class_to_idx = {}
|
||
self.idx_to_class = {}
|
||
|
||
# 加载模型和索引
|
||
self.load_model_and_index()
|
||
|
||
# 创建UI组件
|
||
self.create_widgets()
|
||
|
||
def load_model_and_index(self):
|
||
"""加载训练好的embedding模型和FAISS索引"""
|
||
try:
|
||
# 模型路径
|
||
# 获取当前脚本所在目录
|
||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
# model_path = "../model/embedding_20251011_133653/best_embedding_model.pth"
|
||
model_path = os.path.join(BASE_DIR, "../model/embedding_20251011_133653/best_embedding_model.pth")
|
||
# FAISS索引目录
|
||
# index_dir = "../faiss_vector_db/faiss_index"
|
||
index_dir = os.path.join(BASE_DIR, "../faiss_vector_db/faiss_index")
|
||
|
||
if os.path.exists(model_path) and os.path.exists(index_dir):
|
||
# 1. 加载embedding模型
|
||
print("正在加载embedding模型...")
|
||
# 先把网络结构拿过来，传入embedding维度，预训练，除了加载网络结构，还把预训练的网络参数加载过来
|
||
self.model = create_resnet50_embedding(
|
||
embedding_dim=512,
|
||
pretrained=True,
|
||
use_internal_preprocess=False
|
||
)
|
||
|
||
# 加载模型权重
|
||
checkpoint = torch.load(model_path, map_location=self.device)
|
||
if isinstance(checkpoint, dict):
|
||
if 'model_state_dict' in checkpoint:
|
||
self.model.load_state_dict(checkpoint['model_state_dict'])
|
||
elif 'state_dict' in checkpoint:
|
||
self.model.load_state_dict(checkpoint['state_dict'])
|
||
else:
|
||
self.model.load_state_dict(checkpoint)
|
||
else:
|
||
self.model.load_state_dict(checkpoint)
|
||
# 把模型放到设备上（GPU或CPU）
|
||
self.model.to(self.device)
|
||
# 模型在评估模式下运行时，会禁用某些操作，如BN层和Dropout层
|
||
self.model.eval()
|
||
print("Embedding模型加载成功!")
|
||
|
||
# 2. 加载向量库用于矩阵乘法检索（替代FAISS）
|
||
print("正在加载向量库 embeddings.json ...")
|
||
emb_path = os.path.join(index_dir, 'embeddings.json')
|
||
with open(emb_path, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
# 兼容常见格式：直接列表、字典中的 embeddings/items/vectors、字典列表含 embedding 键
|
||
if isinstance(data, dict):
|
||
vectors = data.get('embeddings') or data.get('items') or data.get('vectors') or []
|
||
if vectors and isinstance(vectors[0], dict) and 'embedding' in vectors[0]:
|
||
vectors = [it['embedding'] for it in vectors]
|
||
else:
|
||
vectors = data
|
||
if not vectors:
|
||
raise RuntimeError("embeddings.json 中未找到向量")
|
||
# 转为 Tensor [N,512]
|
||
self.db_vectors = torch.tensor(vectors, dtype=torch.float32)
|
||
if self.db_vectors.ndim != 2 or self.db_vectors.shape[1] != 512:
|
||
raise RuntimeError(f"embeddings.json 向量形状异常: {tuple(self.db_vectors.shape)}，期望 [N,512]")
|
||
# 预构建转置矩阵 [512,N] 以便快速相乘
|
||
self.db_matrix = self.db_vectors.t().contiguous()
|
||
print(f"向量库已加载: {self.db_vectors.shape[0]} 个向量")
|
||
|
||
# 加载图片路径映射，就是一个list,一张张图片的路径
|
||
paths_path = os.path.join(index_dir, 'image_paths.pkl')
|
||
with open(paths_path, 'rb') as f:
|
||
self.image_paths = pickle.load(f)
|
||
|
||
# 加载标签映射，就是一个list,一张张图片的分类，不过是0，1，2，3，4这种
|
||
labels_path = os.path.join(index_dir, 'labels.pkl')
|
||
with open(labels_path, 'rb') as f:
|
||
self.labels = pickle.load(f)
|
||
|
||
# 加载类别信息
|
||
class_info_path = os.path.join(index_dir, 'class_info.json')
|
||
with open(class_info_path, 'r', encoding='utf-8') as f:
|
||
self.class_info = json.load(f)
|
||
|
||
self.class_names = self.class_info['class_names']
|
||
self.class_to_idx = self.class_info['class_to_idx']
|
||
self.idx_to_class = self.class_info['idx_to_class']
|
||
|
||
print(f"索引元数据已加载: {len(self.image_paths)} 张图片, {len(self.class_names)} 个类别")
|
||
print(f"支持的食物类别: {self.class_names}")
|
||
|
||
else:
|
||
print("模型文件或索引文件不存在，将使用模拟识别")
|
||
self.model = None
|
||
self.db_vectors = None
|
||
self.db_matrix = None
|
||
# 设置默认类别（用于模拟）
|
||
self.class_names = ['回锅肉', '炒细面', '西红柿鸡蛋', '麻辣小面']
|
||
|
||
except Exception as e:
|
||
print(f"模型或索引加载失败: {e}")
|
||
self.model = None
|
||
self.db_vectors = None
|
||
self.db_matrix = None
|
||
# 设置默认类别（用于模拟）
|
||
self.class_names = ['回锅肉', '炒细面', '西红柿鸡蛋', '麻辣小面']
|
||
|
||
def extract_true_class_from_path(self, file_path):
|
||
"""从文件路径中提取真实类别（上一级目录名）"""
|
||
try:
|
||
# 标准化路径
|
||
normalized_path = os.path.normpath(file_path)
|
||
# 获取目录路径
|
||
dir_path = os.path.dirname(normalized_path)
|
||
# 获取上一级目录名（即类别名）
|
||
true_class = os.path.basename(dir_path)
|
||
|
||
# 检查是否是已知的食物类别
|
||
if true_class in self.class_names:
|
||
return true_class
|
||
else:
|
||
# 如果不是已知类别，返回None表示未知
|
||
return None
|
||
|
||
except Exception as e:
|
||
print(f"提取真实类别失败: {e}")
|
||
return None
|
||
|
||
def load_image_with_chinese_path(self, file_path):
|
||
"""使用支持中文路径的方法加载图片"""
|
||
try:
|
||
# 方法1：使用numpy和cv2.imdecode处理中文路径
|
||
with open(file_path, 'rb') as f:
|
||
image_data = f.read()
|
||
|
||
# 将字节数据转换为numpy数组
|
||
nparr = np.frombuffer(image_data, np.uint8)
|
||
|
||
# 使用cv2.imdecode解码图片
|
||
image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
||
|
||
if image is not None:
|
||
return image
|
||
|
||
except Exception as e:
|
||
print(f"加载图片失败: {e}")
|
||
return None
|
||
|
||
def create_widgets(self):
|
||
"""创建UI组件"""
|
||
# 主框架
|
||
self.main_frame = ctk.CTkFrame(self.root)
|
||
self.main_frame.pack(fill="both", expand=True, padx=15, pady=15)
|
||
|
||
# 左侧框架 - 图片上传区域
|
||
self.left_frame = ctk.CTkFrame(self.main_frame, width=600)
|
||
self.left_frame.pack(side="left", fill="both", expand=True, padx=(0, 10), pady=0)
|
||
self.left_frame.pack_propagate(False)
|
||
|
||
# 左侧标题
|
||
self.left_title = ctk.CTkLabel(
|
||
self.left_frame,
|
||
text="图片上传区域 (Embedding相似度检索)",
|
||
font=("Arial", 16, "bold")
|
||
)
|
||
self.left_title.pack(pady=(15, 10))
|
||
|
||
# 拖拽上传区域
|
||
self.upload_frame = ctk.CTkFrame(self.left_frame, fg_color=("gray90", "gray20"))
|
||
self.upload_frame.pack(fill="x", padx=15, pady=(0, 10), ipady=50)
|
||
|
||
# 拖拽提示标签
|
||
self.upload_label = ctk.CTkLabel(
|
||
self.upload_frame,
|
||
text="拖拽图片到这里\n或点击下方按钮选择图片\n支持多图片上传\n基于特征相似度识别",
|
||
font=("Arial", 14),
|
||
text_color=("gray40", "gray60")
|
||
)
|
||
self.upload_label.pack(expand=True)
|
||
|
||
# 绑定拖放事件
|
||
self.upload_frame.drop_target_register(DND_FILES)
|
||
self.upload_frame.dnd_bind('<<Drop>>', self.handle_drop)
|
||
self.upload_frame.bind('<Enter>', self.on_drag_enter)
|
||
self.upload_frame.bind('<Leave>', self.on_drag_leave)
|
||
|
||
# 按钮区域
|
||
self.button_frame = ctk.CTkFrame(self.left_frame)
|
||
self.button_frame.pack(fill="x", padx=15, pady=(0, 10))
|
||
|
||
# 选择图片按钮
|
||
self.select_button = ctk.CTkButton(
|
||
self.button_frame,
|
||
text="选择图片",
|
||
command=self.select_images,
|
||
width=120,
|
||
height=35
|
||
)
|
||
self.select_button.pack(side="left", padx=(10, 5), pady=10)
|
||
|
||
# 清空按钮
|
||
self.clear_button = ctk.CTkButton(
|
||
self.button_frame,
|
||
text="清空图片",
|
||
command=self.clear_images,
|
||
width=120,
|
||
height=35,
|
||
fg_color="gray",
|
||
hover_color="darkgray"
|
||
)
|
||
self.clear_button.pack(side="left", padx=5, pady=10)
|
||
|
||
# 识别按钮
|
||
self.recognize_button = ctk.CTkButton(
|
||
self.button_frame,
|
||
text="开始识别",
|
||
command=self.start_recognition,
|
||
width=120,
|
||
height=35,
|
||
fg_color="green",
|
||
hover_color="darkgreen"
|
||
)
|
||
self.recognize_button.pack(side="right", padx=(5, 10), pady=10)
|
||
self.recognize_button.configure(state="disabled")
|
||
|
||
# 已上传图片显示区域
|
||
self.images_display_frame = ctk.CTkScrollableFrame(
|
||
self.left_frame,
|
||
label_text="已上传的图片"
|
||
)
|
||
self.images_display_frame.pack(fill="both", expand=True, padx=15, pady=(0, 15))
|
||
|
||
# 右侧框架 - 识别结果和新增Embedding区域
|
||
self.right_frame = ctk.CTkFrame(self.main_frame, width=700)
|
||
self.right_frame.pack(side="right", fill="both", expand=True, padx=(10, 0), pady=0)
|
||
self.right_frame.pack_propagate(False)
|
||
|
||
# 创建右侧的选项卡视图
|
||
self.right_tabview = ctk.CTkTabview(self.right_frame)
|
||
self.right_tabview.pack(fill="both", expand=True, padx=15, pady=15)
|
||
|
||
# 识别结果选项卡
|
||
self.results_tab = self.right_tabview.add("识别结果")
|
||
self.create_results_tab()
|
||
|
||
# 新增Embedding选项卡
|
||
self.embedding_tab = self.right_tabview.add("新增Embedding")
|
||
self.create_embedding_tab()
|
||
|
||
# 默认选中识别结果选项卡
|
||
self.right_tabview.set("识别结果")
|
||
|
||
def create_results_tab(self):
|
||
"""创建识别结果选项卡"""
|
||
# 统计信息框架
|
||
self.stats_frame = ctk.CTkFrame(self.results_tab)
|
||
self.stats_frame.pack(fill="x", padx=15, pady=(10, 10))
|
||
|
||
# 统计标签
|
||
self.stats_label = ctk.CTkLabel(
|
||
self.stats_frame,
|
||
text="总图片: 0 | 已识别: 0 | 平均准确率: 0%",
|
||
font=("Arial", 12)
|
||
)
|
||
self.stats_label.pack(pady=10)
|
||
|
||
# 识别结果显示区域
|
||
self.results_display_frame = ctk.CTkScrollableFrame(
|
||
self.results_tab,
|
||
label_text="识别详情 (相似度排序)"
|
||
)
|
||
self.results_display_frame.pack(fill="both", expand=True, padx=15, pady=(0, 15))
|
||
|
||
def create_embedding_tab(self):
|
||
"""创建新增Embedding选项卡"""
|
||
# 新增Embedding区域的变量初始化
|
||
self.embedding_images = [] # 待添加到向量库的图片
|
||
self.selected_class = ctk.StringVar(value="") # 选中的类别
|
||
self.new_class_name = ctk.StringVar(value="") # 新类别名称
|
||
|
||
# 标题
|
||
title_label = ctk.CTkLabel(
|
||
self.embedding_tab,
|
||
text="新增Embedding向量",
|
||
font=("Arial", 16, "bold")
|
||
)
|
||
title_label.pack(pady=(10, 15))
|
||
|
||
# 类别选择区域
|
||
class_frame = ctk.CTkFrame(self.embedding_tab)
|
||
class_frame.pack(fill="x", padx=15, pady=(0, 10))
|
||
|
||
class_title = ctk.CTkLabel(
|
||
class_frame,
|
||
text="选择或新增菜品类别:",
|
||
font=("Arial", 12, "bold")
|
||
)
|
||
class_title.pack(pady=(10, 5))
|
||
|
||
# 现有类别选择
|
||
existing_class_frame = ctk.CTkFrame(class_frame)
|
||
existing_class_frame.pack(fill="x", padx=10, pady=5)
|
||
|
||
existing_label = ctk.CTkLabel(
|
||
existing_class_frame,
|
||
text="选择现有类别:",
|
||
font=("Arial", 11)
|
||
)
|
||
existing_label.pack(side="left", padx=(10, 5), pady=10)
|
||
|
||
self.class_dropdown = ctk.CTkComboBox(
|
||
existing_class_frame,
|
||
values=self.class_names,
|
||
variable=self.selected_class,
|
||
width=200,
|
||
command=self.on_class_selected
|
||
)
|
||
self.class_dropdown.pack(side="left", padx=5, pady=10)
|
||
|
||
# 新类别输入
|
||
new_class_frame = ctk.CTkFrame(class_frame)
|
||
new_class_frame.pack(fill="x", padx=10, pady=5)
|
||
|
||
new_label = ctk.CTkLabel(
|
||
new_class_frame,
|
||
text="或新增类别:",
|
||
font=("Arial", 11)
|
||
)
|
||
new_label.pack(side="left", padx=(10, 5), pady=10)
|
||
|
||
self.new_class_entry = ctk.CTkEntry(
|
||
new_class_frame,
|
||
textvariable=self.new_class_name,
|
||
placeholder_text="输入新的菜品类别名称",
|
||
width=200
|
||
)
|
||
self.new_class_entry.pack(side="left", padx=5, pady=10)
|
||
self.new_class_entry.bind('<KeyRelease>', self.on_new_class_input)
|
||
|
||
# 图片上传区域
|
||
upload_frame = ctk.CTkFrame(self.embedding_tab)
|
||
upload_frame.pack(fill="x", padx=15, pady=(0, 10))
|
||
|
||
upload_title = ctk.CTkLabel(
|
||
upload_frame,
|
||
text="上传样本图片:",
|
||
font=("Arial", 12, "bold")
|
||
)
|
||
upload_title.pack(pady=(10, 5))
|
||
|
||
# 拖拽上传区域
|
||
self.embedding_upload_frame = ctk.CTkFrame(upload_frame, fg_color=("gray90", "gray20"))
|
||
self.embedding_upload_frame.pack(fill="x", padx=10, pady=(0, 10), ipady=30)
|
||
|
||
self.embedding_upload_label = ctk.CTkLabel(
|
||
self.embedding_upload_frame,
|
||
text="拖拽图片到这里或点击下方按钮选择支持多图片上传",
|
||
font=("Arial", 12),
|
||
text_color=("gray40", "gray60")
|
||
)
|
||
self.embedding_upload_label.pack(expand=True)
|
||
|
||
# 绑定拖放事件
|
||
self.embedding_upload_frame.drop_target_register(DND_FILES)
|
||
self.embedding_upload_frame.dnd_bind('<<Drop>>', self.handle_embedding_drop)
|
||
|
||
# 按钮区域
|
||
embedding_button_frame = ctk.CTkFrame(upload_frame)
|
||
embedding_button_frame.pack(fill="x", padx=10, pady=(0, 10))
|
||
|
||
# 选择图片按钮
|
||
self.select_embedding_button = ctk.CTkButton(
|
||
embedding_button_frame,
|
||
text="选择图片",
|
||
command=self.select_embedding_images,
|
||
width=100,
|
||
height=30
|
||
)
|
||
self.select_embedding_button.pack(side="left", padx=(10, 5), pady=5)
|
||
|
||
# 清空图片按钮
|
||
self.clear_embedding_button = ctk.CTkButton(
|
||
embedding_button_frame,
|
||
text="清空图片",
|
||
command=self.clear_embedding_images,
|
||
width=100,
|
||
height=30,
|
||
fg_color="gray",
|
||
hover_color="darkgray"
|
||
)
|
||
self.clear_embedding_button.pack(side="left", padx=5, pady=5)
|
||
|
||
# 添加到向量库按钮
|
||
self.add_embedding_button = ctk.CTkButton(
|
||
embedding_button_frame,
|
||
text="添加到向量库",
|
||
command=self.add_to_vector_db,
|
||
width=120,
|
||
height=30,
|
||
fg_color="green",
|
||
hover_color="darkgreen"
|
||
)
|
||
self.add_embedding_button.pack(side="right", padx=(5, 10), pady=5)
|
||
self.add_embedding_button.configure(state="disabled")
|
||
|
||
# 已选择图片显示区域
|
||
self.embedding_images_frame = ctk.CTkScrollableFrame(
|
||
self.embedding_tab,
|
||
label_text="待添加的图片"
|
||
)
|
||
self.embedding_images_frame.pack(fill="both", expand=True, padx=15, pady=(0, 15))
|
||
|
||
# 状态显示区域
|
||
self.status_frame = ctk.CTkFrame(self.embedding_tab)
|
||
self.status_frame.pack(fill="x", padx=15, pady=(0, 15))
|
||
|
||
self.status_label = ctk.CTkLabel(
|
||
self.status_frame,
|
||
text="请选择类别并上传图片",
|
||
font=("Arial", 11),
|
||
text_color="gray"
|
||
)
|
||
self.status_label.pack(pady=10)
|
||
|
||
def on_class_selected(self, value):
|
||
"""当选择现有类别时"""
|
||
if value:
|
||
self.new_class_name.set("") # 清空新类别输入
|
||
self.update_add_button_state()
|
||
|
||
def on_new_class_input(self, event):
|
||
"""当输入新类别时"""
|
||
if self.new_class_name.get().strip():
|
||
self.selected_class.set("") # 清空现有类别选择
|
||
self.class_dropdown.set("")
|
||
self.update_add_button_state()
|
||
|
||
def update_add_button_state(self):
|
||
"""更新添加按钮状态"""
|
||
has_class = bool(self.selected_class.get() or self.new_class_name.get().strip())
|
||
has_images = bool(self.embedding_images)
|
||
|
||
if has_class and has_images:
|
||
self.add_embedding_button.configure(state="normal")
|
||
if self.new_class_name.get().strip():
|
||
self.status_label.configure(
|
||
text=f"准备为新类别 '{self.new_class_name.get().strip()}' 添加 {len(self.embedding_images)} 张图片",
|
||
text_color="green"
|
||
)
|
||
else:
|
||
self.status_label.configure(
|
||
text=f"准备为类别 '{self.selected_class.get()}' 添加 {len(self.embedding_images)} 张图片",
|
||
text_color="green"
|
||
)
|
||
else:
|
||
self.add_embedding_button.configure(state="disabled")
|
||
if not has_class:
|
||
self.status_label.configure(text="请选择或输入类别", text_color="orange")
|
||
elif not has_images:
|
||
self.status_label.configure(text="请上传图片", text_color="orange")
|
||
|
||
def select_embedding_images(self):
|
||
"""选择要添加到向量库的图片"""
|
||
file_paths = filedialog.askopenfilenames(
|
||
title="选择要添加到向量库的图片",
|
||
filetypes=[
|
||
("图像文件", "*.jpg *.jpeg *.png *.bmp *.gif"),
|
||
("JPEG文件", "*.jpg *.jpeg"),
|
||
("PNG文件", "*.png"),
|
||
("所有文件", "*.*")
|
||
]
|
||
)
|
||
|
||
if file_paths:
|
||
for file_path in file_paths:
|
||
self.add_embedding_image(file_path)
|
||
|
||
def handle_embedding_drop(self, event):
|
||
"""处理拖拽到embedding区域的文件"""
|
||
files = event.data.split()
|
||
for file_path in files:
|
||
# 清理文件路径
|
||
file_path = file_path.strip('{}').strip('"')
|
||
file_path = os.path.normpath(file_path)
|
||
|
||
# 检查是否为图片文件
|
||
valid_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif')
|
||
if file_path.lower().endswith(valid_extensions):
|
||
self.add_embedding_image(file_path)
|
||
|
||
def add_embedding_image(self, file_path):
|
||
"""添加图片到embedding列表"""
|
||
try:
|
||
# 检查文件是否存在
|
||
if not os.path.exists(file_path):
|
||
messagebox.showerror("错误", f"文件不存在: {file_path}")
|
||
return
|
||
|
||
# 检查是否已经添加过
|
||
if file_path in [img['path'] for img in self.embedding_images]:
|
||
messagebox.showinfo("提示", "该图片已经添加过了")
|
||
return
|
||
|
||
# 加载图片
|
||
image = self.load_image_with_chinese_path(file_path)
|
||
if image is None:
|
||
messagebox.showerror("错误", f"无法读取图片: {file_path}")
|
||
return
|
||
|
||
# 添加到列表
|
||
image_info = {
|
||
'path': file_path,
|
||
'name': os.path.basename(file_path),
|
||
'image': image
|
||
}
|
||
self.embedding_images.append(image_info)
|
||
|
||
# 更新显示
|
||
self.update_embedding_images_display()
|
||
self.update_add_button_state()
|
||
|
||
except Exception as e:
|
||
messagebox.showerror("错误", f"添加图片时出错: {str(e)}")
|
||
|
||
def clear_embedding_images(self):
|
||
"""清空embedding图片列表"""
|
||
if self.embedding_images:
|
||
result = messagebox.askyesno("确认", "确定要清空所有待添加的图片吗？")
|
||
if result:
|
||
self.embedding_images.clear()
|
||
self.update_embedding_images_display()
|
||
self.update_add_button_state()
|
||
|
||
def update_embedding_images_display(self):
|
||
"""更新embedding图片显示"""
|
||
# 清空当前显示
|
||
for widget in self.embedding_images_frame.winfo_children():
|
||
widget.destroy()
|
||
|
||
if not self.embedding_images:
|
||
no_image_label = ctk.CTkLabel(
|
||
self.embedding_images_frame,
|
||
text="暂无待添加的图片",
|
||
font=("Arial", 12),
|
||
text_color="gray"
|
||
)
|
||
no_image_label.pack(pady=20)
|
||
return
|
||
|
||
# 显示每张图片
|
||
for i, img_info in enumerate(self.embedding_images):
|
||
# 创建图片框架
|
||
img_frame = ctk.CTkFrame(self.embedding_images_frame)
|
||
img_frame.pack(fill="x", padx=5, pady=5)
|
||
|
||
# 缩放图片用于显示
|
||
display_image = self.resize_image_for_display(img_info['image'], 80, 80)
|
||
display_image = cv2.cvtColor(display_image, cv2.COLOR_BGR2RGB)
|
||
pil_image = Image.fromarray(display_image)
|
||
ctk_image = ctk.CTkImage(light_image=pil_image, dark_image=pil_image, size=(80, 80))
|
||
|
||
# 图片标签
|
||
img_label = ctk.CTkLabel(img_frame, image=ctk_image, text="")
|
||
img_label.image = ctk_image
|
||
img_label.pack(side="left", padx=10, pady=10)
|
||
|
||
# 信息框架
|
||
info_frame = ctk.CTkFrame(img_frame)
|
||
info_frame.pack(side="left", fill="both", expand=True, padx=10, pady=10)
|
||
|
||
# 文件名
|
||
name_label = ctk.CTkLabel(
|
||
info_frame,
|
||
text=f"文件名: {img_info['name']}",
|
||
anchor="w",
|
||
font=("Arial", 11)
|
||
)
|
||
name_label.pack(fill="x", padx=5, pady=2)
|
||
|
||
# 删除按钮
|
||
delete_button = ctk.CTkButton(
|
||
img_frame,
|
||
text="删除",
|
||
command=lambda idx=i: self.remove_embedding_image(idx),
|
||
width=50,
|
||
height=25,
|
||
fg_color="red",
|
||
hover_color="darkred"
|
||
)
|
||
delete_button.pack(side="right", padx=10, pady=10)
|
||
|
||
def remove_embedding_image(self, index):
|
||
"""删除embedding图片"""
|
||
if index < len(self.embedding_images):
|
||
self.embedding_images.pop(index)
|
||
self.update_embedding_images_display()
|
||
self.update_add_button_state()
|
||
|
||
def add_to_vector_db(self):
|
||
"""添加图片到向量数据库"""
|
||
if not self.embedding_images:
|
||
messagebox.showwarning("警告", "请先上传图片")
|
||
return
|
||
|
||
# 确定目标类别
|
||
target_class = self.new_class_name.get().strip() or self.selected_class.get()
|
||
if not target_class:
|
||
messagebox.showwarning("警告", "请选择或输入类别")
|
||
return
|
||
|
||
# 确认操作
|
||
if self.new_class_name.get().strip():
|
||
message = f"确定要创建新类别 '{target_class}' 并添加 {len(self.embedding_images)} 张图片到向量库吗？"
|
||
else:
|
||
message = f"确定要为类别 '{target_class}' 添加 {len(self.embedding_images)} 张图片到向量库吗？"
|
||
|
||
result = messagebox.askyesno("确认", message)
|
||
if not result:
|
||
return
|
||
|
||
# 在新线程中执行添加操作
|
||
self.add_embedding_button.configure(state="disabled", text="添加中...")
|
||
self.status_label.configure(text="正在处理图片并添加到向量库...", text_color="blue")
|
||
|
||
threading.Thread(
|
||
target=self.process_embedding_addition,
|
||
args=(target_class,),
|
||
daemon=True
|
||
).start()
|
||
|
||
def process_embedding_addition(self, target_class):
|
||
"""处理embedding添加的后台任务"""
|
||
try:
|
||
if self.model is None or self.db_vectors is None:
|
||
# 如果模型或向量库未加载，显示模拟结果
|
||
self.root.after(0, lambda: self.show_embedding_result(
|
||
target_class, len(self.embedding_images), True, "模拟模式：图片已添加到向量库"
|
||
))
|
||
return
|
||
|
||
# 处理每张图片
|
||
new_embeddings = []
|
||
new_paths = []
|
||
new_labels = []
|
||
|
||
# 确定类别索引
|
||
if target_class in self.class_names:
|
||
class_idx = self.class_names.index(target_class)
|
||
else:
|
||
# 新类别，添加到类别列表
|
||
class_idx = len(self.class_names)
|
||
self.class_names.append(target_class)
|
||
self.class_to_idx[target_class] = class_idx
|
||
self.idx_to_class[str(class_idx)] = target_class
|
||
|
||
for img_info in self.embedding_images:
|
||
try:
|
||
# 转换图片格式
|
||
image_rgb = cv2.cvtColor(img_info['image'], cv2.COLOR_BGR2RGB)
|
||
pil_image = Image.fromarray(image_rgb)
|
||
|
||
# 提取特征向量
|
||
embedding = self.model.extract_embedding(pil_image, normalize=True)
|
||
embedding = embedding.reshape(1, -1).astype(np.float32)
|
||
|
||
new_embeddings.append(embedding)
|
||
new_paths.append(img_info['path'])
|
||
new_labels.append(class_idx)
|
||
|
||
except Exception as e:
|
||
print(f"处理图片 {img_info['name']} 时出错: {e}")
|
||
continue
|
||
|
||
if new_embeddings:
|
||
# 合并所有新的embedding，并追加到向量库 [N,512]
|
||
all_new_embeddings = np.vstack(new_embeddings).astype(np.float32) # [M,512]
|
||
new_tensor = torch.from_numpy(all_new_embeddings) # [M,512]
|
||
if self.db_vectors is None:
|
||
self.db_vectors = new_tensor
|
||
else:
|
||
self.db_vectors = torch.cat([self.db_vectors, new_tensor], dim=0)
|
||
# 重新构建转置矩阵 [512,N]
|
||
self.db_matrix = self.db_vectors.t().contiguous()
|
||
|
||
# 更新路径和标签列表
|
||
self.image_paths.extend(new_paths)
|
||
self.labels.extend(new_labels)
|
||
|
||
# 保存更新后的索引和元数据
|
||
self.save_updated_index()
|
||
|
||
success_message = f"成功添加 {len(new_embeddings)} 张图片到向量库"
|
||
self.root.after(0, lambda: self.show_embedding_result(
|
||
target_class, len(new_embeddings), True, success_message
|
||
))
|
||
else:
|
||
self.root.after(0, lambda: self.show_embedding_result(
|
||
target_class, 0, False, "没有成功处理任何图片"
|
||
))
|
||
|
||
except Exception as e:
|
||
error_message = f"添加过程中出错: {str(e)}"
|
||
self.root.after(0, lambda: self.show_embedding_result(
|
||
target_class, 0, False, error_message
|
||
))
|
||
|
||
def save_updated_index(self):
|
||
"""保存更新后的索引和元数据"""
|
||
try:
|
||
index_dir = "../faiss_vector_db/faiss_index"
|
||
|
||
# 保存向量库为 embeddings.json（list[list[float]]）
|
||
emb_path = os.path.join(index_dir, 'embeddings.json')
|
||
emb_list = self.db_vectors.cpu().numpy().tolist() if self.db_vectors is not None else []
|
||
with open(emb_path, 'w', encoding='utf-8') as f:
|
||
json.dump(emb_list, f, ensure_ascii=False)
|
||
|
||
# 保存图片路径
|
||
paths_path = os.path.join(index_dir, 'image_paths.pkl')
|
||
with open(paths_path, 'wb') as f:
|
||
# 重新写入图片路径
|
||
pickle.dump(self.image_paths, f)
|
||
|
||
# 保存标签
|
||
labels_path = os.path.join(index_dir, 'labels.pkl')
|
||
with open(labels_path, 'wb') as f:
|
||
# 重新写入标签
|
||
pickle.dump(self.labels, f)
|
||
|
||
# 更新类别信息
|
||
self.class_info = {
|
||
'class_names': self.class_names,
|
||
'class_to_idx': self.class_to_idx,
|
||
'idx_to_class': self.idx_to_class
|
||
}
|
||
|
||
# 保存类别信息
|
||
class_info_path = os.path.join(index_dir, 'class_info.json')
|
||
with open(class_info_path, 'w', encoding='utf-8') as f:
|
||
json.dump(self.class_info, f, ensure_ascii=False, indent=2)
|
||
|
||
print("索引和元数据已成功保存")
|
||
|
||
except Exception as e:
|
||
print(f"保存索引时出错: {e}")
|
||
|
||
def show_embedding_result(self, target_class, count, success, message):
|
||
"""显示embedding添加结果"""
|
||
self.add_embedding_button.configure(state="normal", text="添加到向量库")
|
||
|
||
if success:
|
||
self.status_label.configure(text=message, text_color="green")
|
||
messagebox.showinfo("成功", message)
|
||
|
||
# 清空已添加的图片
|
||
self.embedding_images.clear()
|
||
self.update_embedding_images_display()
|
||
|
||
# 重置选择
|
||
self.selected_class.set("")
|
||
self.new_class_name.set("")
|
||
self.class_dropdown.set("")
|
||
|
||
# 更新类别下拉框（如果有新类别）
|
||
if target_class not in self.class_dropdown.cget("values"):
|
||
current_values = list(self.class_dropdown.cget("values"))
|
||
current_values.append(target_class)
|
||
self.class_dropdown.configure(values=current_values)
|
||
|
||
self.update_add_button_state()
|
||
else:
|
||
self.status_label.configure(text=message, text_color="red")
|
||
messagebox.showerror("错误", message)
|
||
|
||
def select_images(self):
|
||
"""选择图片文件"""
|
||
file_paths = filedialog.askopenfilenames(
|
||
title="选择图片文件",
|
||
filetypes=[
|
||
("图像文件", "*.jpg *.jpeg *.png *.bmp *.gif"),
|
||
("JPEG文件", "*.jpg *.jpeg"),
|
||
("PNG文件", "*.png"),
|
||
("所有文件", "*.*")
|
||
]
|
||
)
|
||
|
||
if file_paths:
|
||
for file_path in file_paths:
|
||
self.add_image(file_path)
|
||
|
||
def handle_drop(self, event):
|
||
"""处理拖拽文件"""
|
||
files = event.data.split()
|
||
for file_path in files:
|
||
# 清理文件路径
|
||
file_path = file_path.strip('{}').strip('"')
|
||
file_path = os.path.normpath(file_path)
|
||
|
||
# 检查是否为图片文件
|
||
valid_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif')
|
||
if file_path.lower().endswith(valid_extensions):
|
||
self.add_image(file_path)
|
||
|
||
def on_drag_enter(self, event):
|
||
"""拖拽进入时的视觉反馈"""
|
||
self.upload_frame.configure(fg_color=("gray80", "gray30"))
|
||
self.upload_label.configure(text="释放鼠标上传图片")
|
||
|
||
def on_drag_leave(self, event):
|
||
"""拖拽离开时恢复正常"""
|
||
self.upload_frame.configure(fg_color=("gray90", "gray20"))
|
||
self.upload_label.configure(text="拖拽图片到这里\n或点击下方按钮选择图片\n支持多图片上传\n基于特征相似度识别")
|
||
|
||
def add_image(self, file_path):
|
||
"""添加图片到上传列表"""
|
||
try:
|
||
# 检查文件是否存在
|
||
if not os.path.exists(file_path):
|
||
messagebox.showerror("错误", f"文件不存在: {file_path}")
|
||
return
|
||
|
||
# 检查是否已经添加过
|
||
if file_path in [img['path'] for img in self.uploaded_images]:
|
||
messagebox.showinfo("提示", "该图片已经添加过了")
|
||
return
|
||
|
||
# 使用支持中文路径的方法加载图片
|
||
image = self.load_image_with_chinese_path(file_path)
|
||
if image is None:
|
||
messagebox.showerror("错误", f"无法读取图片: {file_path}")
|
||
return
|
||
|
||
# 从文件路径中提取真实类别（上一级目录名）
|
||
true_class = self.extract_true_class_from_path(file_path)
|
||
|
||
# 添加到列表
|
||
image_info = {
|
||
'path': file_path,
|
||
'name': os.path.basename(file_path),
|
||
'image': image,
|
||
'true_class': true_class,
|
||
'recognized': False,
|
||
'result': None
|
||
}
|
||
self.uploaded_images.append(image_info)
|
||
|
||
# 更新显示
|
||
self.update_images_display()
|
||
self.update_recognize_button()
|
||
|
||
except Exception as e:
|
||
messagebox.showerror("错误", f"添加图片时出错: {str(e)}")
|
||
|
||
def update_images_display(self):
|
||
"""更新已上传图片的显示"""
|
||
# 清空当前显示
|
||
for widget in self.images_display_frame.winfo_children():
|
||
widget.destroy()
|
||
|
||
# 显示每张图片
|
||
for i, img_info in enumerate(self.uploaded_images):
|
||
# 创建图片框架
|
||
img_frame = ctk.CTkFrame(self.images_display_frame)
|
||
img_frame.pack(fill="x", padx=5, pady=5)
|
||
|
||
# 缩放图片用于显示
|
||
display_image = self.resize_image_for_display(img_info['image'], 100, 100)
|
||
display_image = cv2.cvtColor(display_image, cv2.COLOR_BGR2RGB)
|
||
pil_image = Image.fromarray(display_image)
|
||
ctk_image = ctk.CTkImage(light_image=pil_image, dark_image=pil_image, size=(100, 100))
|
||
|
||
# 图片标签（可点击预览）
|
||
img_label = ctk.CTkLabel(img_frame, image=ctk_image, text="")
|
||
img_label.image = ctk_image # 保持引用
|
||
img_label.pack(side="left", padx=10, pady=10)
|
||
img_label.bind("<Button-1>", lambda e, idx=i: self.preview_image(idx))
|
||
|
||
# 信息框架
|
||
info_frame = ctk.CTkFrame(img_frame)
|
||
info_frame.pack(side="left", fill="both", expand=True, padx=10, pady=10)
|
||
|
||
# 文件名
|
||
name_label = ctk.CTkLabel(
|
||
info_frame,
|
||
text=f"文件名: {img_info['name']}",
|
||
anchor="w"
|
||
)
|
||
name_label.pack(fill="x", padx=5, pady=2)
|
||
|
||
# 状态和真实类别
|
||
if img_info['recognized'] and img_info.get('result'):
|
||
result = img_info['result']
|
||
is_correct = result.get('is_correct')
|
||
if is_correct is True:
|
||
status = "已识别 ✓"
|
||
status_color = "green"
|
||
elif is_correct is False:
|
||
status = "已识别 ✗"
|
||
status_color = "red"
|
||
else:
|
||
status = "已识别 ?"
|
||
status_color = "orange"
|
||
else:
|
||
status = "未识别"
|
||
status_color = None
|
||
|
||
status_label = ctk.CTkLabel(
|
||
info_frame,
|
||
text=f"状态: {status}",
|
||
anchor="w",
|
||
text_color=status_color
|
||
)
|
||
status_label.pack(fill="x", padx=5, pady=2)
|
||
|
||
# 显示真实类别（如果有）
|
||
true_class = img_info.get('true_class')
|
||
if true_class:
|
||
true_class_label = ctk.CTkLabel(
|
||
info_frame,
|
||
text=f"真实类别: {true_class}",
|
||
anchor="w",
|
||
font=("Arial", 10)
|
||
)
|
||
true_class_label.pack(fill="x", padx=5, pady=1)
|
||
|
||
# 删除按钮
|
||
delete_button = ctk.CTkButton(
|
||
img_frame,
|
||
text="删除",
|
||
command=lambda idx=i: self.remove_image(idx),
|
||
width=60,
|
||
height=30,
|
||
fg_color="red",
|
||
hover_color="darkred"
|
||
)
|
||
delete_button.pack(side="right", padx=10, pady=10)
|
||
|
||
def preview_image(self, index):
|
||
"""预览图片"""
|
||
if index >= len(self.uploaded_images):
|
||
return
|
||
|
||
img_info = self.uploaded_images[index]
|
||
|
||
# 创建预览窗口
|
||
preview_window = ctk.CTkToplevel(self.root)
|
||
preview_window.title(f"预览 - {img_info['name']}")
|
||
preview_window.geometry("800x600")
|
||
|
||
# 设置窗口属性，确保在主窗口上方
|
||
preview_window.transient(self.root) # 设置为主窗口的子窗口
|
||
preview_window.grab_set() # 设置为模态窗口
|
||
preview_window.lift() # 提升到最前面
|
||
preview_window.focus_set() # 设置焦点
|
||
|
||
# 居中显示
|
||
preview_window.update_idletasks()
|
||
x = (preview_window.winfo_screenwidth() // 2) - (800 // 2)
|
||
y = (preview_window.winfo_screenheight() // 2) - (600 // 2)
|
||
preview_window.geometry(f"800x600+{x}+{y}")
|
||
|
||
# 显示图片
|
||
display_image = self.resize_image_for_display(img_info['image'], 750, 550)
|
||
display_image = cv2.cvtColor(display_image, cv2.COLOR_BGR2RGB)
|
||
pil_image = Image.fromarray(display_image)
|
||
# 获取实际图片尺寸
|
||
img_width, img_height = pil_image.size
|
||
ctk_image = ctk.CTkImage(light_image=pil_image, dark_image=pil_image, size=(img_width, img_height))
|
||
|
||
img_label = ctk.CTkLabel(preview_window, image=ctk_image, text="")
|
||
img_label.image = ctk_image
|
||
img_label.pack(expand=True, padx=20, pady=20)
|
||
|
||
def remove_image(self, index):
|
||
"""删除图片"""
|
||
if index < len(self.uploaded_images):
|
||
self.uploaded_images.pop(index)
|
||
self.update_images_display()
|
||
self.update_recognize_button()
|
||
self.update_results_display()
|
||
|
||
def clear_images(self):
|
||
"""清空所有图片"""
|
||
if self.uploaded_images:
|
||
result = messagebox.askyesno("确认", "确定要清空所有图片吗？")
|
||
if result:
|
||
self.uploaded_images.clear()
|
||
self.current_results.clear()
|
||
# 重置识别时间
|
||
self.recognition_start_time = None
|
||
self.recognition_duration = 0
|
||
self.update_images_display()
|
||
self.update_recognize_button()
|
||
self.update_results_display()
|
||
self.update_stats()
|
||
|
||
def update_recognize_button(self):
|
||
"""更新识别按钮状态"""
|
||
if self.uploaded_images:
|
||
self.recognize_button.configure(state="normal")
|
||
else:
|
||
self.recognize_button.configure(state="disabled")
|
||
|
||
def start_recognition(self):
|
||
"""开始识别"""
|
||
if not self.uploaded_images:
|
||
messagebox.showinfo("提示", "请先上传图片")
|
||
return
|
||
|
||
# 记录识别开始时间
|
||
self.recognition_start_time = time.time()
|
||
|
||
# 在新线程中执行识别，避免界面卡顿
|
||
self.recognize_button.configure(state="disabled", text="识别中...")
|
||
threading.Thread(target=self.recognize_images, daemon=True).start()
|
||
|
||
def recognize_images(self):
|
||
"""识别所有图片"""
|
||
try:
|
||
self.current_results.clear()
|
||
|
||
for i, img_info in enumerate(self.uploaded_images):
|
||
# 使用embedding相似度识别
|
||
if self.model is not None and self.db_vectors is not None:
|
||
print('矩阵乘法检索')
|
||
# 使用真实的embedding模型和矩阵乘法检索
|
||
predicted_class, confidence, similar_images = self.predict_with_embedding(img_info['image'])
|
||
else:
|
||
print('模拟')
|
||
# 模拟预测结果
|
||
predicted_class = np.random.choice(self.class_names)
|
||
confidence = np.random.uniform(0.6, 0.95)
|
||
similar_images = []
|
||
|
||
# 自动判断识别是否正确
|
||
true_class = img_info.get('true_class')
|
||
is_correct = None
|
||
if true_class is not None:
|
||
is_correct = (predicted_class == true_class)
|
||
|
||
# 保存结果
|
||
result = {
|
||
'image_index': i,
|
||
'image_name': img_info['name'],
|
||
'predicted_class': predicted_class,
|
||
'confidence': confidence,
|
||
'true_class': true_class,
|
||
'is_correct': is_correct,
|
||
'similar_images': similar_images # 相似图片列表
|
||
}
|
||
|
||
self.current_results.append(result)
|
||
img_info['recognized'] = True
|
||
img_info['result'] = result
|
||
|
||
# 更新UI（在主线程中）
|
||
self.root.after(0, self.update_progress, i + 1, len(self.uploaded_images))
|
||
|
||
# 识别完成，更新UI
|
||
self.root.after(0, self.recognition_completed)
|
||
|
||
except Exception as e:
|
||
self.root.after(0, lambda: messagebox.showerror("错误", f"识别过程中出错: {str(e)}"))
|
||
self.root.after(0, self.recognition_completed)
|
||
|
||
def predict_with_embedding(self, image, k=5):
|
||
"""使用embedding向量与向量库矩阵相乘进行Top-K相似检索并投票"""
|
||
try:
|
||
# 将OpenCV图像转换为PIL图像
|
||
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||
pil_image = Image.fromarray(image_rgb)
|
||
|
||
# 提取查询图片的特征向量 [1,512]
|
||
query_embedding = self.model.extract_embedding(pil_image, normalize=True).astype(np.float32)
|
||
query_tensor = torch.from_numpy(query_embedding.reshape(1, -1)) # [1,512]
|
||
# 打印query_tensor的前十个数字
|
||
|
||
# print(query_embedding[:50])
|
||
|
||
if self.db_matrix is None:
|
||
raise RuntimeError("向量库未加载")
|
||
|
||
# 点积相似度： [1,512] x [512,N] -> [1,N]
|
||
sims = torch.matmul(query_tensor, self.db_matrix) # [1,N]
|
||
|
||
# Top-K
|
||
topk_scores, topk_indices = torch.topk(sims, k=min(k, sims.shape[1]), dim=1)
|
||
indices = topk_indices.cpu().numpy()
|
||
scores = topk_scores.cpu().numpy()
|
||
print("相似度索引：", indices)
|
||
print("相似度分数：", scores)
|
||
|
||
# 收集相似图片的类别
|
||
similar_classes = []
|
||
similar_images = []
|
||
|
||
for i in range(indices.shape[1]):
|
||
idx = int(indices[0][i])
|
||
if idx < len(self.labels):
|
||
score = float(scores[0][i])
|
||
class_idx = self.labels[idx]
|
||
class_name = self.class_names[class_idx] if 0 <= class_idx < len(self.class_names) else str(class_idx)
|
||
image_path = self.image_paths[idx] if 0 <= idx < len(self.image_paths) else ""
|
||
|
||
similar_classes.append(class_name)
|
||
similar_images.append({
|
||
'path': image_path,
|
||
'class': class_name,
|
||
'score': score
|
||
})
|
||
|
||
# 使用投票机制确定最终预测类别
|
||
if similar_classes:
|
||
class_counts = Counter(similar_classes)
|
||
predicted_class = class_counts.most_common(1)[0][0]
|
||
|
||
max_score = float(scores[0][0]) if scores.shape[1] > 0 else 0.0
|
||
vote_ratio = class_counts[predicted_class] / len(similar_classes)
|
||
confidence = max_score * vote_ratio
|
||
|
||
return predicted_class, confidence, similar_images
|
||
else:
|
||
predicted_class = np.random.choice(self.class_names)
|
||
confidence = 0.1
|
||
return predicted_class, confidence, []
|
||
|
||
except Exception as e:
|
||
print(f"Embedding预测出错: {e}")
|
||
predicted_class = np.random.choice(self.class_names)
|
||
confidence = np.random.uniform(0.1, 0.3)
|
||
return predicted_class, confidence, []
|
||
|
||
def update_progress(self, current, total):
|
||
"""更新识别进度"""
|
||
self.recognize_button.configure(text=f"识别中... ({current}/{total})")
|
||
self.update_images_display()
|
||
self.update_results_display()
|
||
|
||
def recognition_completed(self):
|
||
"""识别完成"""
|
||
# 计算识别耗时
|
||
if self.recognition_start_time is not None:
|
||
self.recognition_duration = time.time() - self.recognition_start_time
|
||
|
||
self.recognize_button.configure(state="normal", text="开始识别")
|
||
self.update_stats()
|
||
messagebox.showinfo("完成", f"所有图片识别完成！识别耗时: {self.recognition_duration:.2f}秒")
|
||
|
||
def resize_image_for_display(self, image, max_width, max_height):
|
||
"""调整图片大小用于显示"""
|
||
height, width = image.shape[:2]
|
||
scale = min(max_width / width, max_height / height)
|
||
|
||
if scale < 1:
|
||
new_width = int(width * scale)
|
||
new_height = int(height * scale)
|
||
return cv2.resize(image, (new_width, new_height))
|
||
|
||
return image
|
||
|
||
def update_results_display(self):
|
||
"""更新识别结果显示"""
|
||
# 清空当前显示
|
||
for widget in self.results_display_frame.winfo_children():
|
||
widget.destroy()
|
||
|
||
if not self.current_results:
|
||
no_result_label = ctk.CTkLabel(
|
||
self.results_display_frame,
|
||
text="暂无识别结果",
|
||
font=("Arial", 14),
|
||
text_color="gray"
|
||
)
|
||
no_result_label.pack(pady=20)
|
||
return
|
||
|
||
# 显示每个识别结果
|
||
for i, result in enumerate(self.current_results):
|
||
# 结果框架
|
||
result_frame = ctk.CTkFrame(self.results_display_frame)
|
||
result_frame.pack(fill="x", padx=5, pady=5)
|
||
|
||
# 获取原始图片
|
||
img_info = self.uploaded_images[result['image_index']]
|
||
display_image = self.resize_image_for_display(img_info['image'], 120, 120)
|
||
display_image = cv2.cvtColor(display_image, cv2.COLOR_BGR2RGB)
|
||
pil_image = Image.fromarray(display_image)
|
||
ctk_image = ctk.CTkImage(light_image=pil_image, dark_image=pil_image, size=(120, 120))
|
||
|
||
# 图片标签
|
||
img_label = ctk.CTkLabel(result_frame, image=ctk_image, text="")
|
||
img_label.image = ctk_image
|
||
img_label.pack(side="left", padx=10, pady=10)
|
||
img_label.bind("<Button-1>", lambda e, idx=result['image_index']: self.preview_image(idx))
|
||
|
||
# 信息框架
|
||
info_frame = ctk.CTkFrame(result_frame)
|
||
info_frame.pack(side="left", fill="both", expand=True, padx=10, pady=10)
|
||
|
||
# 文件名
|
||
name_label = ctk.CTkLabel(
|
||
info_frame,
|
||
text=f"文件: {result['image_name']}",
|
||
anchor="w",
|
||
font=("Arial", 12, "bold")
|
||
)
|
||
name_label.pack(fill="x", padx=5, pady=2)
|
||
|
||
# 识别结果标题
|
||
result_title_label = ctk.CTkLabel(
|
||
info_frame,
|
||
text="识别结果 (基于相似度):",
|
||
anchor="w",
|
||
font=("Arial", 11)
|
||
)
|
||
result_title_label.pack(fill="x", padx=5, pady=(2, 0))
|
||
|
||
# 识别结果内容（大字体、加粗，颜色根据正确性决定）
|
||
is_correct = result.get('is_correct')
|
||
if is_correct is True:
|
||
result_color = "green" # 识别正确显示绿色
|
||
elif is_correct is False:
|
||
result_color = "red" # 识别错误显示红色
|
||
else:
|
||
result_color = "orange" # 无法判断显示橙色
|
||
|
||
result_content_label = ctk.CTkLabel(
|
||
info_frame,
|
||
text=result['predicted_class'],
|
||
anchor="w",
|
||
font=("Arial", 18, "bold"),
|
||
text_color=result_color
|
||
)
|
||
result_content_label.pack(fill="x", padx=5, pady=(0, 2))
|
||
|
||
# 置信度
|
||
confidence_label = ctk.CTkLabel(
|
||
info_frame,
|
||
text=f"相似度得分: {result['confidence']:.3f}",
|
||
anchor="w",
|
||
font=("Arial", 11)
|
||
)
|
||
confidence_label.pack(fill="x", padx=5, pady=2)
|
||
|
||
# 真实类别
|
||
true_class = result.get('true_class')
|
||
if true_class is not None:
|
||
true_class_label = ctk.CTkLabel(
|
||
info_frame,
|
||
text=f"真实类别: {true_class}",
|
||
anchor="w",
|
||
font=("Arial", 11)
|
||
)
|
||
true_class_label.pack(fill="x", padx=5, pady=2)
|
||
|
||
# 自动判断结果
|
||
is_correct = result.get('is_correct')
|
||
if is_correct is not None:
|
||
if is_correct:
|
||
status_text = "✓ 识别正确"
|
||
status_color = "green"
|
||
else:
|
||
status_text = "✗ 识别错误"
|
||
status_color = "red"
|
||
|
||
status_label = ctk.CTkLabel(
|
||
info_frame,
|
||
text=status_text,
|
||
anchor="w",
|
||
font=("Arial", 12, "bold"),
|
||
text_color=status_color
|
||
)
|
||
status_label.pack(fill="x", padx=5, pady=5)
|
||
else:
|
||
# 如果无法自动判断，显示未知状态
|
||
status_label = ctk.CTkLabel(
|
||
info_frame,
|
||
text="? 无法自动判断（路径中未包含已知类别）",
|
||
anchor="w",
|
||
font=("Arial", 11),
|
||
text_color="orange"
|
||
)
|
||
status_label.pack(fill="x", padx=5, pady=5)
|
||
|
||
# 显示相似图片信息（如果有）
|
||
similar_images = result.get('similar_images', [])
|
||
if similar_images:
|
||
similar_label = ctk.CTkLabel(
|
||
info_frame,
|
||
text=f"基于前{len(similar_images)}张最相似图片的投票结果",
|
||
anchor="w",
|
||
font=("Arial", 10),
|
||
text_color="gray"
|
||
)
|
||
similar_label.pack(fill="x", padx=5, pady=(5, 2))
|
||
|
||
# 显示前5张最相似的图片信息
|
||
for j, sim_img in enumerate(similar_images[:5]):
|
||
# 创建一个水平框架来放置信息和按钮
|
||
sim_frame = ctk.CTkFrame(info_frame)
|
||
sim_frame.pack(fill="x", padx=15, pady=1)
|
||
|
||
# 相似图片信息
|
||
sim_info = f" {j+1}. {sim_img['class']} (相似度: {sim_img['score']:.3f})"
|
||
sim_info_label = ctk.CTkLabel(
|
||
sim_frame,
|
||
text=sim_info,
|
||
anchor="w",
|
||
font=("Arial", 9),
|
||
text_color="gray"
|
||
)
|
||
sim_info_label.pack(side="left", fill="x", expand=True, padx=5, pady=2)
|
||
|
||
# 查看按钮（使用眼睛图标）
|
||
view_btn = ctk.CTkButton(
|
||
sim_frame,
|
||
text="👁", # 眼睛图标
|
||
width=25,
|
||
height=20,
|
||
font=("Arial", 12),
|
||
command=lambda path=sim_img['path']: self.show_similar_image(path),
|
||
fg_color="#2196F3",
|
||
hover_color="#1976D2",
|
||
corner_radius=10
|
||
)
|
||
view_btn.pack(side="right", padx=5, pady=2)
|
||
|
||
def show_similar_image(self, image_path):
|
||
"""显示相似图片的窗口"""
|
||
try:
|
||
# 创建新窗口
|
||
similar_window = ctk.CTkToplevel(self.root)
|
||
similar_window.title("相似图片查看")
|
||
similar_window.geometry("600x500")
|
||
similar_window.resizable(True, True)
|
||
|
||
# 设置窗口属性，确保在主窗口上方
|
||
similar_window.transient(self.root) # 设置为主窗口的子窗口
|
||
similar_window.grab_set() # 设置为模态窗口
|
||
similar_window.lift() # 提升到最前面
|
||
similar_window.focus_set() # 设置焦点
|
||
|
||
# 居中显示
|
||
similar_window.update_idletasks()
|
||
x = (similar_window.winfo_screenwidth() // 2) - (600 // 2)
|
||
y = (similar_window.winfo_screenheight() // 2) - (500 // 2)
|
||
similar_window.geometry(f"600x500+{x}+{y}")
|
||
|
||
# 主框架
|
||
main_frame = ctk.CTkFrame(similar_window)
|
||
main_frame.pack(fill="both", expand=True, padx=10, pady=10)
|
||
|
||
# 标题
|
||
title_label = ctk.CTkLabel(
|
||
main_frame,
|
||
text="相似图片详情",
|
||
font=("Arial", 16, "bold")
|
||
)
|
||
title_label.pack(pady=(10, 20))
|
||
|
||
# 图片路径信息
|
||
path_label = ctk.CTkLabel(
|
||
main_frame,
|
||
text=f"图片路径: {image_path}",
|
||
font=("Arial", 10),
|
||
text_color="gray",
|
||
wraplength=550
|
||
)
|
||
path_label.pack(pady=(0, 10))
|
||
|
||
# 图片显示框架
|
||
image_frame = ctk.CTkFrame(main_frame)
|
||
image_frame.pack(fill="both", expand=True, padx=10, pady=10)
|
||
|
||
# 加载并显示图片
|
||
try:
|
||
# 打开图片
|
||
pil_image = Image.open(image_path)
|
||
|
||
# 计算合适的显示尺寸
|
||
max_width, max_height = 500, 350
|
||
img_width, img_height = pil_image.size
|
||
|
||
# 计算缩放比例
|
||
scale = min(max_width / img_width, max_height / img_height, 1.0)
|
||
new_width = int(img_width * scale)
|
||
new_height = int(img_height * scale)
|
||
|
||
# 调整图片大小
|
||
pil_image = pil_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||
|
||
# 转换为CTk图片
|
||
ctk_image = ctk.CTkImage(
|
||
light_image=pil_image,
|
||
dark_image=pil_image,
|
||
size=(new_width, new_height)
|
||
)
|
||
|
||
# 显示图片
|
||
image_label = ctk.CTkLabel(
|
||
image_frame,
|
||
image=ctk_image,
|
||
text=""
|
||
)
|
||
image_label.pack(expand=True, pady=20)
|
||
|
||
# 图片信息
|
||
info_text = f"原始尺寸: {img_width} × {img_height}显示尺寸: {new_width} × {new_height}"
|
||
info_label = ctk.CTkLabel(
|
||
main_frame,
|
||
text=info_text,
|
||
font=("Arial", 9),
|
||
text_color="gray"
|
||
)
|
||
info_label.pack(pady=(0, 10))
|
||
|
||
except Exception as e:
|
||
error_label = ctk.CTkLabel(
|
||
image_frame,
|
||
text=f"无法加载图片: {str(e)}",
|
||
font=("Arial", 12),
|
||
text_color="red"
|
||
)
|
||
error_label.pack(expand=True)
|
||
|
||
# 关闭按钮
|
||
close_btn = ctk.CTkButton(
|
||
main_frame,
|
||
text="关闭",
|
||
width=100,
|
||
command=similar_window.destroy,
|
||
fg_color="#f44336",
|
||
hover_color="#da190b"
|
||
)
|
||
close_btn.pack(pady=(10, 10))
|
||
|
||
# 让窗口获得焦点
|
||
similar_window.focus()
|
||
similar_window.lift()
|
||
|
||
except Exception as e:
|
||
messagebox.showerror("错误", f"显示相似图片时出错: {str(e)}")
|
||
|
||
def update_stats(self):
|
||
"""更新统计信息"""
|
||
total_images = len(self.uploaded_images)
|
||
recognized_images = len(self.current_results)
|
||
|
||
# 计算准确率（基于自动判断的结果）
|
||
auto_judged_results = [r for r in self.current_results if r.get('is_correct') is not None]
|
||
if auto_judged_results:
|
||
correct_count = sum(1 for r in auto_judged_results if r['is_correct'])
|
||
accuracy = (correct_count / len(auto_judged_results)) * 100
|
||
|
||
# 构建统计文本，包含识别时间
|
||
if self.recognition_duration > 0:
|
||
stats_text = f"总图片: {total_images} | 已识别: {recognized_images} | 可判断: {len(auto_judged_results)} | 识别时间: {self.recognition_duration:.2f}秒 | 平均准确率: {accuracy:.1f}%"
|
||
else:
|
||
stats_text = f"总图片: {total_images} | 已识别: {recognized_images} | 可判断: {len(auto_judged_results)} | 平均准确率: {accuracy:.1f}%"
|
||
else:
|
||
# 如果没有可自动判断的结果
|
||
if self.recognition_duration > 0:
|
||
stats_text = f"总图片: {total_images} | 已识别: {recognized_images} | 可判断: 0 | 识别时间: {self.recognition_duration:.2f}秒 | 平均准确率: 0%"
|
||
else:
|
||
stats_text = f"总图片: {total_images} | 已识别: {recognized_images} | 可判断: 0 | 平均准确率: 0%"
|
||
|
||
self.stats_label.configure(text=stats_text)
|
||
|
||
|
||
def main():
|
||
# 创建支持拖放的窗口
|
||
root = TkinterDnD.Tk()
|
||
app = EmbeddingFoodClassifierApp(root)
|
||
root.mainloop()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |