树莓派OPENCV数字识别——中
首先,先说一个我试了很久很久的结论,那就是支持向量机这种方法效果太差了,我最开始想着使用支持向量机做,因为他最简单,但是实际效果差得离谱,最后还是选择了用pytorch搞深度学习。对于是否需要优化模型,我这里并不需要,因为我的准确率已经高到99%了,我会继续尝试别的训练数字,看会不会有影响,有影响了这里我再改。中篇主要是关于模型训练部分的,上篇已经生成了一个文件夹的数字,那么基于那个数字集,进行训
·
首先,先说一个我试了很久很久的结论,那就是支持向量机这种方法效果太差了,我最开始想着使用支持向量机做,因为他最简单,但是实际效果差得离谱,最后还是选择了用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不是一般的麻烦。
火山引擎开发者社区是火山引擎打造的AI技术生态平台,聚焦Agent与大模型开发,提供豆包系列模型(图像/视频/视觉)、智能分析与会话工具,并配套评测集、动手实验室及行业案例库。社区通过技术沙龙、挑战赛等活动促进开发者成长,新用户可领50万Tokens权益,助力构建智能应用。
更多推荐
所有评论(0)