yixiong/FoodClassifier
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

目前可以正常识别了，支持多张图片同时识别！

Browse Source
This commit is contained in:
zhanghuan
2025-09-03 11:11:44 +08:00
parent a60cd58a0b
commit 0496ada9f0
3 changed files with 824 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
classifier/food_classifier_app.py
+691
View File
@@ -0,0 +1,691 @@
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
# 设置customtkinter的外观
ctk.set_appearance_mode("System")
ctk.set_default_color_theme("blue")
# 定义CNN模型（与训练代码中的结构相同）
class FoodCNN(nn.Module):
def __init__(self):
super(FoodCNN, self).__init__()
# 第一个卷积块
self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
self.conv2 = nn.Conv2d(32, 32, 3, padding=1)
self.pool1 = nn.MaxPool2d(2, 2)
self.dropout1 = nn.Dropout2d(0.25)
# 第二个卷积块
self.conv3 = nn.Conv2d(32, 64, 3, padding=1)
self.conv4 = nn.Conv2d(64, 64, 3, padding=1)
self.pool2 = nn.MaxPool2d(2, 2)
self.dropout2 = nn.Dropout2d(0.25)
# 第三个卷积块
self.conv5 = nn.Conv2d(64, 128, 3, padding=1)
self.conv6 = nn.Conv2d(128, 128, 3, padding=1)
self.pool3 = nn.MaxPool2d(2, 2)
self.dropout3 = nn.Dropout2d(0.25)
# 全连接层
self.fc1 = nn.Linear(128 * 4 * 4, 512)
self.dropout4 = nn.Dropout(0.5)
self.fc2 = nn.Linear(512, 2) # 2分类
def forward(self, x):
# 第一个卷积块
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = self.pool1(x)
x = self.dropout1(x)
# 第二个卷积块
x = F.relu(self.conv3(x))
x = F.relu(self.conv4(x))
x = self.pool2(x)
x = self.dropout2(x)
# 第三个卷积块
x = F.relu(self.conv5(x))
x = F.relu(self.conv6(x))
x = self.pool3(x)
x = self.dropout3(x)
# 展平
x = x.view(-1, 128 * 4 * 4)
# 全连接层
x = F.relu(self.fc1(x))
x = self.dropout4(x)
x = self.fc2(x)
return x
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.model = None
self.transform = None
self.load_model()
# 创建UI组件
self.create_widgets()
def load_model(self):
"""加载训练好的PyTorch模型"""
try:
model_path = "../model/01/best_food_model.pth"
if os.path.exists(model_path):
# 创建模型实例
self.model = FoodCNN()
# 加载模型权重
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 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使用RGB，OpenCV使用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
# 添加到列表
image_info = {
'path': file_path,
'name': os.path.basename(file_path),
'image': image,
'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)
# 状态
status = "已识别" if img_info['recognized'] else "未识别"
status_label = ctk.CTkLabel(
info_frame,
text=f"状态: {status}",
anchor="w"
)
status_label.pack(fill="x", padx=5, pady=2)
# 删除按钮
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")
# 显示图片
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.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.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)
# 保存结果
result = {
'image_index': i,
'image_name': img_info['name'],
'predicted_class': predicted_class,
'confidence': confidence,
'is_correct': None # 用户可以手动标记
}
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):
"""识别完成"""
self.recognize_button.configure(state="normal", text="开始识别")
self.update_stats()
messagebox.showinfo("完成", "所有图片识别完成！")
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_label = ctk.CTkLabel(
info_frame,
text=f"识别结果: {result['predicted_class']}",
anchor="w",
font=("Arial", 11)
)
result_label.pack(fill="x", padx=5, pady=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)
# 正确性标记框架
correct_frame = ctk.CTkFrame(info_frame)
correct_frame.pack(fill="x", padx=5, pady=5)
correct_label = ctk.CTkLabel(correct_frame, text="识别是否正确:")
correct_label.pack(side="left", padx=5)
# 正确按钮
correct_button = ctk.CTkButton(
correct_frame,
text="正确",
command=lambda idx=i: self.mark_result(idx, True),
width=60,
height=25,
fg_color="green" if result['is_correct'] is True else "gray",
hover_color="darkgreen"
)
correct_button.pack(side="left", padx=2)
# 错误按钮
incorrect_button = ctk.CTkButton(
correct_frame,
text="错误",
command=lambda idx=i: self.mark_result(idx, False),
width=60,
height=25,
fg_color="red" if result['is_correct'] is False else "gray",
hover_color="darkred"
)
incorrect_button.pack(side="left", padx=2)
def mark_result(self, result_index, is_correct):
"""标记识别结果的正确性"""
if result_index < len(self.current_results):
self.current_results[result_index]['is_correct'] = is_correct
self.update_results_display()
self.update_stats()
def update_stats(self):
"""更新统计信息"""
total_images = len(self.uploaded_images)
recognized_images = len(self.current_results)
# 计算准确率
marked_results = [r for r in self.current_results if r['is_correct'] is not None]
if marked_results:
correct_count = sum(1 for r in marked_results if r['is_correct'])
accuracy = (correct_count / len(marked_results)) * 100
else:
accuracy = 0
# 更新统计标签
stats_text = f"总图片: {total_images} | 已识别: {recognized_images} | 平均准确率: {accuracy:.1f}%"
self.stats_label.configure(text=stats_text)
def main():
# 创建支持拖放的窗口
root = TkinterDnD.Tk()
app = FoodClassifierApp(root)
root.mainloop()
if __name__ == "__main__":
main()