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

任务描述

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

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

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


代码介绍

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

cat_dog_test.py代码:

python 复制代码
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代码:

python 复制代码
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)),  # 调整图片尺寸为224x224
    transforms.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),输出通道为32
        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)  # 输出为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))  # 全连接层 + ReLU
        x = 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.0
    correct = 0
    total = 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


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

相关推荐
小沛94 分钟前
Spring AI Alibaba JManus:前后端异步消息回显机制深度解析
人工智能·spring·jmanus
Allen_LVyingbo5 分钟前
智能AI医疗物资/耗材管理系统升级改造方案分析
人工智能·自然语言处理·开源·健康医疗·扩展屏应用开发
图灵学术计算机期刊速推8 分钟前
AAAI‘26 | 聚焦人工智能前沿:西工大李学龙教授荣任赞助主席,论文取号逼近三万,精彩不容错过!
人工智能·机器学习·计算机视觉
wa的一声哭了23 分钟前
Python多进程并行multiprocess基础
开发语言·jvm·人工智能·python·机器学习·语言模型·自然语言处理
全宝25 分钟前
🎨【AI绘画实战】从零搭建Stable Diffusion环境,手把手教你生成超可爱Q版大头照!
人工智能·python·stable diffusion
chy存钱罐32 分钟前
数据集:机器学习的基石
人工智能·机器学习·scikit-learn
Ice__Cai36 分钟前
Django 视图详解(View):处理请求与返回响应的核心
数据库·后端·python·django·pip
枫叶丹437 分钟前
【Qt开发】信号与槽(三)-> 自定义信号和槽
开发语言·qt
蔷薇のぬ43 分钟前
Python 使用pandas库实现Excel字典码表对照自动化处理
python·pandas
青衫客361 小时前
LLM—— 基于 MCP 协议(SSE 模式)的工具调用实践
python·大模型·llm