手写数字识别----实时识别

训练模型

我们选择用深度学习训练手写数字识别。主要是因为深度学习模型能够自动学习数据中的特征，而不需要手动设计特征提取器。对于手写数字识别，底层可能学到笔画的基本特征，而高层次则可能学到数字的整体形状。可以使用在大规模数据上预训练的模型，在相对较小的手写数字数据集上进行微调。

#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

# 定义手写数字识别模型
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)

        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)

        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.fc2(x)

        return x

# 下载并准备数据集
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 初始化模型、损失函数和优化器
model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 20

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0

    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / len(train_loader)}")

# 保存模型
torch.save(model.state_dict(), 'mnist_model.pth')

print("Model saved.")

利用模型进行实时推理

import cv2
import torch
from torchvision import transforms
from PIL import Image
from train import SimpleNet

# 加载预训练的 PyTorch 模型
pytorch_model = SimpleNet()
pytorch_model.load_state_dict(torch.load('mnist_model.pth'))
pytorch_model.eval()

# 设置摄像头
cap = cv2.VideoCapture(0)  # 0 表示默认摄像头

# 图像预处理
transform = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

while True:
    # 从摄像头读取一帧
    ret, frame = cap.read()

    # 转为灰度图
    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # 转为 PIL 图像
    pil_image = Image.fromarray(gray_frame)

    # 图像预处理
    input_image = transform(pil_image)
    input_image = input_image.unsqueeze(0)  

    # 使用 PyTorch 模型进行推理
    with torch.no_grad():
        output = pytorch_model(input_image)

    # 获取预测结果
    _, predicted = torch.max(output, 1)

    # 在图像上显示识别结果
    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(frame, f"预测的数字: {predicted.item()}", (10, 50), font, 1, (255, 0, 0), 2, cv2.LINE_AA)

    # 显示实时图像
    cv2.imshow('实时手写数字识别', frame)

    # 按 'q' 键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像头并关闭窗口
cap.release()
cv2.destroyAllWindows()

结论

当然，它还能做些更有意思的事情，下期做手写电话号码识别。