当前位置：首页 > news >正文

基于卷积神经网络和Pyqt5的猫狗识别小程序

news 来源：原创 2025/5/10 7:13:15

任务描述

猫狗分类任务（Dogs vs Cats）是Kaggle平台在2013年举办的一个经典计算机视觉竞赛。官方给出的Kaggle Dogs vs Cats 数据集中包括由12500张猫咪图片和12500张狗狗图片组成的训练集，12500张未标记照片组成的测试集。选手需要在规定时间内搭建模型，并在测试集中取得尽可能高的准确率。

数据集中的图片长这样👆，是不是还蛮贴近生活呢~

本文将设计一个简单的三层卷积神经网络，完成分类任务。基于Pyqt5设计交互界面，在测试集中检验分类结果。工作量不大非常适合新手作为入门项目~最终功能实现如下图所示。

代码介绍

我觉得新手最有成就感的事情莫过于跑通一个功能了，我们抛开原理不谈，先完整运行程序。代码我是用pycharm写的，新建一个classfy_program项目，在classfy_program项目中创建两个py文件，分别用于训练和测试数据集。此外还创建data文件夹，用于存放训练和测试用的图片。

cat_dog_test.py代码：

import sys
import torch
import torch.nn as nn
from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QPushButton, QVBoxLayout, QFileDialog
from PyQt5.QtGui import QPixmap
from PyQt5.QtCore import Qt
from torchvision import transforms
from PIL import Image# ---------------- 定义模型结构 ----------------
class CNN(nn.Module):def __init__(self):super(CNN, self).__init__()self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)self.pool = nn.MaxPool2d(kernel_size=2, stride=2)self.fc1 = nn.Linear(128 * 28 * 28, 512)self.fc2 = nn.Linear(512, 2)  # 猫 vs 狗def forward(self, x):x = self.pool(torch.relu(self.conv1(x)))x = self.pool(torch.relu(self.conv2(x)))x = self.pool(torch.relu(self.conv3(x)))x = x.view(-1, 128 * 28 * 28)x = torch.relu(self.fc1(x))x = self.fc2(x)return x# ---------------- 图像预处理 ----------------
transform = transforms.Compose([transforms.Resize((224, 224)),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]),
])# 类别名称（根据你的训练数据文件夹顺序）
class_names = ['猫', '狗']# ---------------- 主窗口 ----------------
class CatDogClassifier(QWidget):def __init__(self):super().__init__()self.setWindowTitle("猫狗分类器")self.setGeometry(100, 100, 400, 500)# 加载模型self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")self.model = CNN().to(self.device)self.model.load_state_dict(torch.load("model_checkpoint.pth", map_location=self.device))self.model.eval()# 创建界面组件self.image_label = QLabel("请点击按钮加载图片", self)self.image_label.setAlignment(Qt.AlignCenter)self.image_label.setFixedSize(300, 300)self.result_label = QLabel("", self)self.result_label.setAlignment(Qt.AlignCenter)self.load_button = QPushButton("加载图片", self)self.load_button.clicked.connect(self.load_image)# 布局layout = QVBoxLayout()layout.addWidget(self.image_label)layout.addWidget(self.result_label)layout.addWidget(self.load_button)self.setLayout(layout)def load_image(self):file_path, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "Image Files (*.png *.jpg *.jpeg)")if file_path:# 显示图片pixmap = QPixmap(file_path)pixmap = pixmap.scaled(300, 300, Qt.KeepAspectRatio)self.image_label.setPixmap(pixmap)# 进行预测image = Image.open(file_path).convert("RGB")input_tensor = transform(image).unsqueeze(0).to(self.device)with torch.no_grad():output = self.model(input_tensor)_, predicted = torch.max(output, 1)result = class_names[predicted.item()]self.result_label.setText(f"预测结果：{result}")# ---------------- 运行主程序 ----------------
if __name__ == "__main__":app = QApplication(sys.argv)window = CatDogClassifier()window.show()sys.exit(app.exec_())

cat_dog_train.py代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms# 数据预处理与加载
transform = transforms.Compose([transforms.Resize((224,224)),  # 调整图片尺寸为224x224transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 适用于自定义模型的标准化
])# 设置训练集和验证集
train_dataset = datasets.ImageFolder(root='data/train', transform=transform)train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)# 自定义卷积神经网络模型
class CNN(nn.Module):def __init__(self):super(CNN, self).__init__()# 第一层卷积self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)  # 输入通道为3（RGB），输出通道为32self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)# 池化层self.pool = nn.MaxPool2d(kernel_size=2, stride=2)# 全连接层self.fc1 = nn.Linear(128 * 28 * 28, 512)  # 调整为经过池化后的维度self.fc2 = nn.Linear(512, 2)  # 输出为2（猫 vs 狗）def forward(self, x):x = self.pool(torch.relu(self.conv1(x)))  # 第一层卷积 + ReLU + 最大池化x = self.pool(torch.relu(self.conv2(x)))  # 第二层卷积 + ReLU + 最大池化x = self.pool(torch.relu(self.conv3(x)))  # 第三层卷积 + ReLU + 最大池化x = x.view(-1, 128 * 28 * 28)  # 展平x = torch.relu(self.fc1(x))  # 全连接层 + ReLUx = self.fc2(x)  # 输出层return x# 实例化模型
model = CNN()# 使用GPU训练（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)  # 将模型移动到GPU（如果可用）# 损失函数与优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.0001)# 训练模型
num_epochs = 10
for epoch in range(num_epochs):model.train()running_loss = 0.0correct = 0total = 0# 训练阶段for images, labels in train_loader:images, labels = images.to(device), labels.to(device)  # 将数据移动到GPU（如果可用）optimizer.zero_grad()outputs = model(images)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()_, predicted = torch.max(outputs, 1)total += labels.size(0)correct += (predicted == labels).sum().item()# 输出每轮训练损失与准确率print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(train_loader):.4f}, Accuracy: {100 * correct / total:.2f}%")# 保存模型检查点torch.save(model.state_dict(), 'model_checkpoint.pth')print("Training complete!")

data数据集的下载链接在这里了：

Download Kaggle Cats and Dogs Dataset from Official Microsoft Download Center

希望看到这里的大家能够点一个小小的赞❤❤，后续会持续更新更多内容~~