首先,先说一个我试了很久很久的结论,那就是支持向量机这种方法效果太差了,我最开始想着使用支持向量机做,因为他最简单,但是实际效果差得离谱,最后还是选择了用pytorch搞深度学习。

中篇主要是关于模型训练部分的,上篇已经生成了一个文件夹的数字,那么基于那个数字集,进行训练集的生成。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from PIL import Image
import os


# 定义一个简单的CNN模型用于数字识别
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, 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.fc1 = nn.Linear(64 * 16 * 16, 128)  # 输入特征的大小需要根据图像尺寸来计算
        self.fc2 = nn.Linear(128, 10)  # 10个输出类别,对应数字0-9

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = x.view(x.size(0), -1)  # 展平
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


# 定义图像变换
transform = transforms.Compose([
    transforms.Resize((64, 64)),  # 确保图像大小为64x64
    transforms.ToTensor(),  # 转换为Tensor
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),  # 归一化
])


# 自定义数据集类
class DigitDataset(torch.utils.data.Dataset):
    def __init__(self, image_folder, transform=None):
        self.image_folder = image_folder
        self.transform = transform
        self.image_paths = []
        self.labels = []
        self._load_images()

    def _load_images(self):
        for filename in os.listdir(self.image_folder):
            if filename.endswith(".png"):
                self.image_paths.append(os.path.join(self.image_folder, filename))
                label = int(filename.split('_')[1].split('.')[0])  # 从文件名中提取标签
                self.labels.append(label)

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = Image.open(img_path).convert('RGB')
        label = self.labels[idx]

        if self.transform:
            image = self.transform(image)

        return image, label


# 加载数据集
dataset = DigitDataset(image_folder='generated_images', transform=transform)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

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

# 训练模型
epochs = 5
for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    for images, labels in dataloader:
        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}/{epochs}], Loss: {running_loss / len(dataloader):.4f}, Accuracy: {100 * correct / total:.2f}%')

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

这里有一点声明一下,要下什么库这点,请自行查询,因为不同python版本的函数不太一样,根据自己灵活调整。

对于是否需要优化模型,我这里并不需要,因为我的准确率已经高到99%了,我会继续尝试别的训练数字,看会不会有影响,有影响了这里我再改。

生成的这个训练结果发送给树莓派即可。

下篇为数字识别的代码

如果有时间我会把支持向量机的做法也写了,因为树莓派部署pytorch不是一般的麻烦。

Logo

火山引擎开发者社区是火山引擎打造的AI技术生态平台,聚焦Agent与大模型开发,提供豆包系列模型(图像/视频/视觉)、智能分析与会话工具,并配套评测集、动手实验室及行业案例库。社区通过技术沙龙、挑战赛等活动促进开发者成长,新用户可领50万Tokens权益,助力构建智能应用。

更多推荐