Files
FoodClassifier/classifier/food_classifier_app.py
T

756 lines
29 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import os
import json
import cv2
import numpy as np
import customtkinter as ctk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from typing import List, Optional, Tuple
from tkinterdnd2 import DND_FILES, TkinterDnD
import threading
import time
from net import create_food_cnn
from settings import settings
# 设置customtkinter的外观
ctk.set_appearance_mode("System")
ctk.set_default_color_theme("blue")
class FoodClassifierApp:
def __init__(self, root):
self.root = root
self.root.title("数字味道-食物识别系统")
self.root.geometry("1400x800")
# 食物类别(根据您的数据集)
self.food_classes = ["回锅肉", "西红柿鸡蛋","麻辣小面"]
# 设备设置
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
# 模型相关
self.model = None
self.transform = None
self.load_model()
# 创建UI组件
self.create_widgets()
def load_model(self):
"""加载训练好的PyTorch模型"""
try:
model_path = settings.INFERENCE_BEST_MODEL_PATH
if os.path.exists(model_path):
# 创建模型实例
self.model = create_food_cnn()
# 加载模型权重
self.model.load_state_dict(torch.load(model_path, map_location=self.device))
self.model.to(self.device)
self.model.eval() # 设置为评估模式
# 定义图像预处理(与训练时相同)
self.transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
print("PyTorch模型加载成功")
else:
print("模型文件不存在,将使用模拟识别")
self.model = None
except Exception as e:
print(f"模型加载失败: {e}")
self.model = None
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.food_classes:
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
# 方法2:如果方法1失败,尝试使用PIL
from PIL import Image as PILImage
pil_image = PILImage.open(file_path)
# 转换为RGB(如果是RGBA
if pil_image.mode == 'RGBA':
pil_image = pil_image.convert('RGB')
elif pil_image.mode == 'L': # 灰度图
pil_image = pil_image.convert('RGB')
# 转换为numpy数组
image_array = np.array(pil_image)
# PIL使用RGBOpenCV使用BGR,需要转换
image = cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR)
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="图片上传区域",
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支持多图片上传",
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))
# 右侧框架 - 识别结果区域
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_title = ctk.CTkLabel(
self.right_frame,
text="识别结果",
font=("Arial", 16, "bold")
)
self.right_title.pack(pady=(15, 10))
# 统计信息框架
self.stats_frame = ctk.CTkFrame(self.right_frame)
self.stats_frame.pack(fill="x", padx=15, pady=(0, 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.right_frame,
label_text="识别详情"
)
self.results_display_frame.pack(fill="both", expand=True, padx=15, pady=(0, 15))
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支持多图片上传")
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)
tk_image = ImageTk.PhotoImage(pil_image)
# 图片标签(可点击预览)
img_label = ctk.CTkLabel(img_frame, image=tk_image, text="")
img_label.image = tk_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)
tk_image = ImageTk.PhotoImage(pil_image)
img_label = ctk.CTkLabel(preview_window, image=tk_image, text="")
img_label.image = tk_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):
# 预处理图片
if self.model is not None and self.transform is not None:
# 使用真实模型预测
prediction, confidence = self.predict_with_model(img_info['image'])
predicted_class = self.food_classes[prediction]
else:
# 模拟预测结果
prediction = np.random.randint(0, len(self.food_classes))
predicted_class = self.food_classes[prediction]
confidence = np.random.uniform(0.6, 0.95)
# 自动判断识别是否正确
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 # 自动判断结果
}
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_model(self, image):
"""使用PyTorch模型进行预测"""
try:
# 将OpenCV图像转换为PIL图像
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
pil_image = Image.fromarray(image_rgb)
# 应用预处理
input_tensor = self.transform(pil_image).unsqueeze(0) # 添加batch维度
input_tensor = input_tensor.to(self.device)
# 进行预测
with torch.no_grad():
outputs = self.model(input_tensor)
probabilities = F.softmax(outputs, dim=1)
confidence, predicted = torch.max(probabilities, 1)
return predicted.item(), confidence.item()
except Exception as e:
print(f"模型预测出错: {e}")
# 返回随机结果作为备选
prediction = np.random.randint(0, len(self.food_classes))
confidence = np.random.uniform(0.6, 0.95)
return prediction, 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)
tk_image = ImageTk.PhotoImage(pil_image)
# 图片标签
img_label = ctk.CTkLabel(result_frame, image=tk_image, text="")
img_label.image = tk_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']:.2%}",
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)
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 = FoodClassifierApp(root)
root.mainloop()
if __name__ == "__main__":
main()