Yanick's Blog

import torchvision
import matplotlib.pyplot as plt
import numpy as np

# 加载原始 MNIST 训练集（无 transform，保留原始格式）
train_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=True,
    download=True,
)

# 获取前2个样本的原始数据和标签
data = train_dataset.data  # 形状 (60000, 28, 28)，uint8 0-255
targets = train_dataset.targets  # 形状 (60000,)

# 提取前2个样本
img1_tensor = data[0]  # 第1个样本：张量 (28, 28)
label1 = targets[0].item()  # 第1个样本的标签（转成Python数字）
img2_tensor = data[1]  # 第2个样本：张量 (28, 28)
label2 = targets[1].item()  # 第2个样本的标签

# 格式转换：PyTorch 张量 → NumPy 数组（matplotlib 支持 NumPy 格式）
img1_np = img1_tensor.numpy()
img2_np = img2_tensor.numpy()

# 可视化前2个样本（并排显示）
plt.figure(figsize=(8, 4))

# 第1个样本
plt.subplot(1, 2, 1)
plt.imshow(img1_np, cmap='gray')  # cmap='gray' 表示灰度图显示
plt.title(f'训练集第1个样本\n真实标签: {label1}')
plt.axis('off')  # 隐藏坐标轴，更干净

# 第2个样本
plt.subplot(1, 2, 2)
plt.imshow(img2_np, cmap='gray')
plt.title(f'训练集第2个样本\n真实标签: {label2}')
plt.axis('off')

plt.tight_layout()  # 自动调整子图间距
plt.show()

# 打印原始数据的关键信息（帮助理解）
print("="*50)
print("训练集第1个样本信息：")
print(f"张量形状: {img1_tensor.shape} → (高, 宽)（无通道维度，原始灰度图）")
print(f"数据类型: {img1_tensor.dtype} → uint8（0-255 整数像素值）")
print(f"像素值范围: [{img1_tensor.min()}, {img1_tensor.max()}]")
print(f"真实标签: {label1}")

print("\n训练集第2个样本信息：")
print(f"张量形状: {img2_tensor.shape}")
print(f"像素值范围: [{img2_tensor.min()}, {img2_tensor.max()}]")
print(f"真实标签: {label2}")
print("="*50)

import torch
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose([
    transforms.ToTensor(),  # 核心：图片 → Tensor
])

# 2. 加载训练集（download=True 首次运行会自动下载）
train_dataset = torchvision.datasets.MNIST(
    root='./data',  # 数据保存路径
    train=True,     # True=训练集，False=测试集
    download=True,
    transform=transform
)

# 3. 加载测试集
test_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=False,
    download=True,
    transform=transform
)

# 4. 批量加载数据（DataLoader 自动打乱、分批，提升训练效率）
batch_size = 64  # 每批处理 64 张图片
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True  # 训练集打乱，增加泛化性
)
test_loader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=batch_size,
    shuffle=False  # 测试集无需打乱
)

# 验证数据：查看一批数据的形状
images, labels = next(iter(train_loader))
print("单批图片形状:", images.shape)  # 输出 (64, 1, 28, 28) → (批量数, 通道数, 高, 宽)
print("单批标签形状:", labels.shape)  # 输

import torch
import torch.nn as nn
import torch.nn.functional as F

class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        # 输入层：28×28 图片 → 展平为 784 维向量
        self.fc1 = nn.Linear(28*28, 256)  # 第一层全连接：784 → 256（隐藏层神经元数）
        self.fc2 = nn.Linear(256, 128)    # 第二层全连接：256 → 128
        self.fc3 = nn.Linear(128, 10)     # 输出层：128 → 10（对应 0-9 共 10 类）

    def forward(self, x):
        # 前向传播：定义数据流动路径
        x = x.view(-1, 28*28)  # 展平：(batch, 1, 28, 28) → (batch, 784)
        x = F.relu(self.fc1(x)) # 激活函数 ReLU（引入非线性）
        x = F.relu(self.fc2(x))
        x = self.fc3(x)         # 输出层无需激活（后续用交叉熵损失，内置 Softmax）
        return x

# 初始化模型
model = MLP()
# 移到 GPU（如果可用）
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
model.to(device)
print(model)  # 打印模型结构，确认无误

import torch.optim as optim

# 1. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 分类问题专用损失（内置 Softmax）
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)  # SGD 优化器，学习率 0.01

# 2. 训练循环（epoch：遍历整个训练集的次数，这里设 5 次）
num_epochs = 5
for epoch in range(num_epochs):
    model.train()  # 切换到训练模式（启用 Dropout/BatchNorm 等）
    running_loss = 0.0  # 记录当前 epoch 的总损失

    # 遍历训练集的每一批数据
    for i, (images, labels) in enumerate(train_loader, 0):
        # 数据移到 GPU/CPU
        images, labels = images.to(device), labels.to(device)

        # 关键步骤：
        optimizer.zero_grad()  # 清空上一批的梯度（避免累积）
        outputs = model(images)  # 前向传播：模型预测
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播：计算梯度
        optimizer.step()  # 更新参数：梯度下降

        # 打印训练进度
        running_loss += loss.item()
        if i % 100 == 99:  # 每 100 批打印一次
            print(f'Epoch [{epoch+1}/{num_epochs}], Batch [{i+1}/{len(train_loader)}], Loss: {running_loss/100:.4f}')
            running_loss = 0.0

print('训练完成！')

# 保存模型
torch.save(model.state_dict(), 'digit_recognition_model.pth')
# 加载模型
model.load_state_dict(torch.load('digit_recognition_model.pth'))

import tkinter as tk
from tkinter import messagebox
from PIL import Image, ImageDraw
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

# -------------------------- 1. 配置参数（需根据你的训练情况修改）--------------------------
# 模型路径（替换为你训练时保存的模型路径）
MODEL_PATH = "digit_recognition_model.pth"
# 模型类型（选择你训练的模型：MLP 或 CNN）
MODEL_TYPE = "MLP"  # 可选 "MLP" 或 "CNN"
# 设备配置（M 系列 Mac 用 MPS，其他用 CPU）
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")


# -------------------------- 2. 定义模型结构（必须和训练时完全一致！）--------------------------
# 多层感知机（MLP）
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 28 * 28)  # 展平
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)


# -------------------------- 3. 加载训练好的模型--------------------------
def load_model():
    # 初始化模型（和训练时一致）
    if MODEL_TYPE == "MLP":
        model = MLP().to(device)
    else:
        raise ValueError("MODEL_TYPE 只能是 'MLP'")

    # 加载模型参数
    try:
        model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
        model.eval()  # 切换到评估模式
        print(f"模型加载成功！使用设备：{device}")
        return model
    except FileNotFoundError:
        messagebox.showerror("错误", f"未找到模型文件：{MODEL_PATH}\n请先训练模型并保存，或修改 MODEL_PATH")
        exit()


# 加载模型（程序启动时执行）
model = load_model()


# -------------------------- 4. 图像预处理函数（关键：和训练时一致）--------------------------
def preprocess_image(image):
    """
    输入：PIL 图像（300×300 白底黑字）
    输出：模型可接收的 Tensor（(1, 1, 28, 28)）
    """
    # 定义预处理流程
    transform_list = [
        transforms.Resize((28, 28)),  # 缩放到 28×28
        transforms.ToTensor(),  # 转 Tensor（0-255 → 0-1）
    ]

    transform = transforms.Compose(transform_list)

    # 预处理步骤
    img = image.convert("L")  # 转灰度图
    img = img.point(lambda x: 255 - x)  # 反色：白底黑字 → 黑底白字（匹配 MNIST 格式）
    img_tensor = transform(img).unsqueeze(0)  # 添加 batch 维度 (1, 1, 28, 28)
    return img_tensor.to(device), img  # 返回 Tensor 和预处理后的 PIL 图（用于可视化）


# -------------------------- 5. 模型推理函数--------------------------
def predict_digit(image):
    """输入手写图像，返回预测结果"""
    img_tensor, processed_img = preprocess_image(image)

    with torch.no_grad():  # 关闭梯度计算
        output = model(img_tensor)
        pred_label = torch.max(output, 1)[1].item()  # 取概率最大的标签
        pred_prob = torch.softmax(output, 1)[0][pred_label].item()  # 计算预测概率

    return pred_label, pred_prob, processed_img


# -------------------------- 6. Tkinter GUI 手写画布--------------------------
class HandwritingApp:
    def __init__(self, root):
        self.root = root
        self.root.title("手写数字识别（28×28 MNIST 风格）")
        self.root.geometry("400x400")  # 窗口大小

        # 画布配置（300×300，白色背景）
        self.canvas = tk.Canvas(root, width=300, height=300, bg="white", bd=2, relief=tk.SUNKEN)
        self.canvas.pack(pady=20)

        # 按钮配置
        self.btn_frame = tk.Frame(root)
        self.btn_frame.pack()

        self.predict_btn = tk.Button(self.btn_frame, text="识别数字", command=self.on_predict, font=("Arial", 12))
        self.predict_btn.grid(row=0, column=0, padx=10)

        self.clear_btn = tk.Button(self.btn_frame, text="清空画布", command=self.on_clear, font=("Arial", 12))
        self.clear_btn.grid(row=0, column=1, padx=10)

        # 绘制状态
        self.drawing = False
        self.last_x = None
        self.last_y = None

        # 创建 PIL 图像和绘图对象（用于保存画布内容）
        self.pil_image = Image.new("RGB", (300, 300), "white")
        self.draw = ImageDraw.Draw(self.pil_image)

        # 绑定鼠标事件
        self.canvas.bind("<Button-1>", self.start_drawing)  # 鼠标按下
        self.canvas.bind("<B1-Motion>", self.draw_line)  # 鼠标拖动
        self.canvas.bind("<ButtonRelease-1>", self.stop_drawing)  # 鼠标释放

    def start_drawing(self, event):
        """开始绘制（鼠标按下）"""
        self.drawing = True
        self.last_x = event.x
        self.last_y = event.y

    def draw_line(self, event):
        """绘制线条（鼠标拖动）"""
        if self.drawing:
            # 在 Tkinter 画布上画（可视化）
            self.canvas.create_line(
                self.last_x, self.last_y, event.x, event.y,
                fill="black", width=15, capstyle=tk.ROUND, smooth=tk.TRUE
            )
            # 在 PIL 图像上画（用于后续保存和预处理）
            self.draw.line(
                [(self.last_x, self.last_y), (event.x, event.y)],
                fill="black", width=15, joint="round"
            )
            # 更新最后坐标
            self.last_x = event.x
            self.last_y = event.y

    def stop_drawing(self, event):
        """停止绘制（鼠标释放）"""
        self.drawing = False

    def on_clear(self, ):
        """清空画布"""
        self.canvas.delete("all")  # 清空 Tkinter 画布
        self.pil_image = Image.new("RGB", (300, 300), "white")  # 重置 PIL 图像
        self.draw = ImageDraw.Draw(self.pil_image)

    def on_predict(self):
        """识别数字（点击按钮触发）"""
        try:
            # 推理手写数字
            pred_label, pred_prob, processed_img = predict_digit(self.pil_image)

            # 显示结果（弹窗 + 可视化预处理后的 28×28 图）
            messagebox.showinfo(
                "预测结果",
                f"预测数字：{pred_label}\n预测概率：{pred_prob:.2%}"
            )

            # 可视化预处理后的 28×28 图（和 MNIST 格式一致）
            plt.figure(figsize=(4, 4))
            plt.imshow(processed_img, cmap="gray")
            plt.title(f"预处理后（28×28）→ 预测：{pred_label}")
            plt.axis("off")
            plt.show()

        except Exception as e:
            messagebox.showerror("错误", f"识别失败：{str(e)}")


# -------------------------- 7. 运行 GUI 程序--------------------------
if __name__ == "__main__":
    root = tk.Tk()
    app = HandwritingApp(root)
    root.mainloop()

import torch
import torch.nn as nn
import torch.nn.functional as F


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 卷积层 1：输入 1 通道 → 32 通道，卷积核 3×3，步长 1， padding 1（保持尺寸不变）
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        # 池化层 1：最大池化 2×2（尺寸减半：28→14）
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        # 卷积层 2：32 通道 → 64 通道
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        # 全连接层 1：64 通道×7×7（池化后尺寸）→ 128 维
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  # conv1 → ReLU → pool
        x = self.pool(F.relu(self.conv2(x)))  # conv2 → ReLU → pool（尺寸再减半：14→7）
        x = x.view(-1, 64 * 7 * 7)  # 展平：(batch, 64, 7, 7) → (batch, 64*7*7)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST 标准归一化
])