造车先做三蹦子-之三:自制数据集(6x6数据集)230103

6*6的数据集制造、与识别：

#6*6的数据集的制作、与识别、测试、输出等
 
 
import torch
import torch.nn as nn
import torch.optim as optim
 
# 定义模型
class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, num_classes)
        self.relu = nn.ReLU()
 
    def forward(self, x):
        out = self.relu(self.fc1(x))
        out = self.relu(self.fc2(out))
        out = self.fc3(out)
        return out
 
# 数据准备
train_data = torch.tensor([
    [[0,0,1,0,0,0],[0,0,1,0,0,0],[0,0,1,0,0,0],[0,0,1,0,0,0],[0,0,1,0,0,0],[0,0,1,0,0,0]], 
    [[0,0,0,0,0,0],[0,0,0,0,0,0],[1,1,1,1,1,1],[1,1,1,1,1,1],[0,0,0,0,0,0],[0,0,0,0,0,0]],
    # ... 其他训练数据
]  , dtype=torch.float32 )
train_labels = torch.tensor([
    [1,0,0,0,0,0,0],
    [0,1,0,0,0,0,0],
    # ... 其他训练标签
]  , dtype=torch.float32 )
 
test_data = torch.tensor([
    [[0,0,0,0,0,0],[0,1,1,1,1,0],[1,1,1,1,1,1],[0,0,0,0,0,0],[0,0,0,0,0,0],[0,0,0,0,0,0]],
    [[0,0,1,0,0,0],[0,0,1,0,0,0],[0,0,1,0,0,0],[0,0,1,0,0,0],[0,0,1,0,0,0],[0,0,0,1,0,0]],
    # ... 其他测试数据
]  , dtype=torch.float32 )
 
# 参数设置
input_size = 6 * 6  # 输入维度
hidden_size = 50   # 隐藏层节点数量
num_classes = 7    # 输出维度
learning_rate = 0.001
num_epochs = 1000
 
model = NeuralNet(input_size, hidden_size, num_classes)
 
# 损失和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
 
# 训练模型
for epoch in range(num_epochs):
    # 前向传播
    outputs = model(train_data.view(-1, 36))
    loss = criterion(outputs, torch.argmax(train_labels, 1))
    
    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
 
    if (epoch+1) % 100 == 0:
        print ('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))
 
# 保存模型
torch.save(model.state_dict(), 'model.ckpt')
 
# 测试
with torch.no_grad():
    test_outputs = model(test_data.view(-1, 36))
    predicted = torch.argmax(test_outputs.data, 1)
    for i, prediction in enumerate(predicted):
        print(f"Test data {i+1}: Predicted label: {prediction.item()}")

#true Type 字体作为训练库，完成字符的识别
 
import torch
import torch.nn as nn
import torch.optim as optim
from PIL import Image, ImageDraw, ImageFont
from torchvision import transforms
import matplotlib.pyplot as plt
 
Times2000=1000
 
# 参数设置
font_path = "arial.ttf"    #e:\\arial.ttf"
#siz28e28 = 28
siz28e28=28
characters = ["2","4"] +[str(i) for i in range(8,10)] + ["A","Z"]   #["A", "B"]
print(characters)
 
# 1. 生成字符和数字的点阵数据
def render_char(char, siz28e28):
    image = Image.new("L", (siz28e28, siz28e28), "white")
    draw = ImageDraw.Draw(image)
    font = ImageFont.truetype(font_path, siz28e28)
#-----------------------------------------
    w03 = draw.textlength(char, font=font)
    h03 = siz28e28
    print("{[w3",w03, "h3",w03 ,"]} ")
#-----------------------------------------
    draw.text(((siz28e28-w03)/2, (siz28e28-h03)/2), char, font=font, fill="black")
    return image
 
data = []
labels = []
for i, char in enumerate(characters):
    img = render_char(char, siz28e28)
    data.append(transforms.ToTensor()(img))
    labels.append(i)
    
# 2. 训练神经网络模型
class SimpleNet(nn.Module):
    def __init__(self, num_classes):
        super(SimpleNet, self).__init__()
        self.fc = nn.Linear(siz28e28 * siz28e28, num_classes)
 
    def forward(self, x):
        x = x.view(-1, siz28e28 * siz28e28)
        x = self.fc(x)
        return x
    
model = SimpleNet(len(characters))
loss_function = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
 
for epoch in range(Times2000):#8000):#1000):
    inputs = torch.stack(data)
    targets = torch.tensor(labels)
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()
 
# 3. 使用模型进行预测
def predict_image(img_path):
    model.eval()
    img = Image.open(img_path).convert("L").resize((siz28e28, siz28e28))
    img_tensor = transforms.ToTensor()(img).unsqueeze(0)
    output = model(img_tensor)
    _, predicted = output.max(1)
    return characters[predicted[0]]
 
 
#//1-----
im="f:\\22letter23r1002\8_16x16.png"
predicted_char = predict_image(im)
print(f"预测的字符为: {predicted_char}")
 
# 使用matplotlib显示结果
plt.imshow(Image.open(im))
#plt.title(f"Predicted: {predicted_char}")
#plt.axis("off")
#plt.show()