diff --git a/classifier/embedding_food_classifier_search_app.py b/classifier/embedding_food_classifier_search_app.py new file mode 100644 index 0000000..5287db1 --- /dev/null +++ b/classifier/embedding_food_classifier_search_app.py @@ -0,0 +1,1527 @@ +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: + # 模型路径 + model_path = "../model/embedding_20250930_102826/best_embedding_model.pth" + + # FAISS索引目录 + index_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('<>', self.handle_drop) + self.upload_frame.bind('', self.on_drag_enter) + self.upload_frame.bind('', 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('', 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('<>', 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("", 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[:10]) + + 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("", 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() \ No newline at end of file diff --git a/net/resnet_embedding.py b/net/resnet_embedding.py index af0b938..4516fca 100644 --- a/net/resnet_embedding.py +++ b/net/resnet_embedding.py @@ -97,21 +97,31 @@ class ResNet50EmbeddingNet(nn.Module): def mobile_preprocess(self, x: torch.Tensor) -> torch.Tensor: """ 移动端预处理（TorchScript兼容） - + Args: x: 输入tensor，形状为 [batch, 3, height, width]，值范围 0-1 - + Returns: torch.Tensor: 预处理后的tensor，形状为 [batch, 3, 224, 224] """ + # # 通过 scale_factor 避免对 size 的整数检查（规避 RecursionError） + # h, w = x.shape[2], x.shape[3] + # # 防御：避免除零 + # if h == 0 or w == 0: + # raise ValueError(f"Invalid input size: height={h}, width={w}") + # scale_h = 224.0 / float(h) + # scale_w = 224.0 / float(w) + # + # x = F.interpolate(x, scale_factor=(scale_h, scale_w), + # mode='bilinear', align_corners=False) # 缩放到 224x224 x = F.interpolate(x, size=(224, 224), mode='bilinear', align_corners=False) - + # ImageNet标准化 mean = torch.tensor([0.485, 0.456, 0.406], device=x.device).view(1, 3, 1, 1) std = torch.tensor([0.229, 0.224, 0.225], device=x.device).view(1, 3, 1, 1) x = (x - mean) / std - + return x def _forward_backbone(self, x: torch.Tensor) -> torch.Tensor: @@ -180,8 +190,8 @@ class ResNet50EmbeddingNet(nn.Module): # 移动端预处理 x = self.mobile_preprocess(x) - # 执行前向传播 - return self.forward(x, normalize=normalize) + # 显式调用基类 forward，避免子类重写的 forward 形成递归 + return ResNet50EmbeddingNet.forward(self, x, normalize=normalize) def extract_embedding(self, image: Union[Image.Image, np.ndarray, torch.Tensor], normalize: bool = True) -> np.ndarray: @@ -210,7 +220,8 @@ class ResNet50EmbeddingNet(nn.Module): # 提取特征 embedding = self.forward(x, normalize=normalize) - + # embedding = self.forward(x, normalize=False) + return embedding.cpu().numpy().flatten() def extract_batch_embeddings(self, images: List[Union[Image.Image, np.ndarray]], @@ -365,7 +376,7 @@ def create_resnet50_embedding(embedding_dim: int = 512, ) -def create_mobile_resnet50_embedding(embedding_dim: int = 512) -> ResNet50EmbeddingNet: +def create_mobile_resnet50_embedding(embedding_dim: int = 512,pretrained: bool = True) -> ResNet50EmbeddingNet: """ 创建移动端ResNet50 Embedding模型 @@ -387,15 +398,10 @@ def create_mobile_resnet50_embedding(embedding_dim: int = 512) -> ResNet50Embedd def forward(self, x: torch.Tensor, normalize: bool = True) -> torch.Tensor: """ 移动端前向传播（自动包含预处理） - - Args: - x: 输入tensor，形状为 [batch, 3, height, width]，值范围 0-1 - normalize: 是否对输出进行L2归一化 - - Returns: - torch.Tensor: embedding向量 """ - return self.forward_mobile(x, normalize=normalize) + # 在子类里做预处理，然后显式调用基类 forward + x = self.mobile_preprocess(x) + return ResNet50EmbeddingNet.forward(self, x, normalize=normalize) return MobileResNet50Embedding(embedding_dim=embedding_dim) diff --git a/toAndroid/toAndroidEmbedding.py b/toAndroid/toAndroidEmbedding.py index 0a5ed58..79bb16a 100644 --- a/toAndroid/toAndroidEmbedding.py +++ b/toAndroid/toAndroidEmbedding.py @@ -15,7 +15,8 @@ def main(): print("开始转换ResNet50 Embedding模型为移动端格式...") # 1. 加载训练好的embedding模型权重 - base_model = create_resnet50_embedding(embedding_dim=512, pretrained=True) + # base_model = create_resnet50_embedding(embedding_dim=512, pretrained=True) + base_model = create_mobile_resnet50_embedding(embedding_dim=512, pretrained=True) model_path = "../model/embedding_20250917_145342/best_embedding_model.pth" if not os.path.exists(model_path): diff --git a/toAndroid/toAndroidSearch.py b/toAndroid/toAndroidSearch.py index 204157b..f65b604 100644 --- a/toAndroid/toAndroidSearch.py +++ b/toAndroid/toAndroidSearch.py @@ -127,11 +127,11 @@ class TwoLayerSimilarityNet(nn.Module): raise ValueError("输入张量形状必须为 [batch, 512]") # 输入归一化 - # x_norm = F.normalize(x, p=2, dim=1) # [B, 512] + x_norm = F.normalize(x, p=2, dim=1) # [B, 512] # 单层线性变换 - # sims = self.layer2(x_norm) # [B, N] - sims = self.layer2(x) # [B, N] + sims = self.layer2(x_norm) # [B, N] + # sims = self.layer2(x) # [B, N] return sims