import os
import numpy as np
import torch
import torch.nn as nn
import torch.functional as F
from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image
from torch.utils.data import DataLoader
import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.autograd import Variable
captcha_array = list("0123456789abcdefghijklmnopqrstuvwxyz")
def text2vec(x):
zeros = torch.zeros((4, 36), dtype = torch.long)
for i in range(len(x)):
zeros[i, captcha_array.index(x[i])] = 1
vec = zeros
return vec
def vec2text(x):
label = ''
label_tensor = torch.argmax(x, dim=1)
for i in label_tensor:
label += captcha_array[i.data]
return label
def single_vec_2_text(label_tensor):
label = ''
for i in label_tensor:
label += captcha_array[i.data]
return label
def batch_vec_2_text(x):
label_batch = []
label_all = x.argmax(dim=2)
for index_image in range(label_all.size(0)):
label_batch.append(single_vec_2_text(label_all[index_image]))
return label_batch
def compare_list(x, y):
result_list = []
for le in range(len(x)):
if x[le] == y[le]:
result_list.append(True)
return result_list.count(True), len(x), result_list.count(True)/len(x)
class MyNet(nn.Module):
def __init__(self):
super(MyNet, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 32, (3, 3), (1, 1), 1), # in torch.Size([batch, 1, 60, 160])
# nn.BatchNorm2d(32),
nn.ReLU(inplace = True),
nn.MaxPool2d(2) # out (batch, 32, 30, 80)
)
self.layer2 = nn.Sequential(
nn.Conv2d(32, 64, (3, 3), (1, 1), 1), # in torch.Size([
# nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(2) # out (batch, 64, 15, 40)
)
self.layer3 = nn.Sequential(
nn.Conv2d(64, 128, (3, 3), (1, 1), 1), # in torch.Size([
# nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.MaxPool2d(2) # out (batch, 128, 7, 20)
)
self.layer4 = nn.Sequential(
nn.Conv2d(128, 256, (3, 3), (1, 1), 1), # in torch.Size(
# nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
nn.MaxPool2d(2) # out (batch, 256, 3, 10)
)
self.layer5 = nn.Sequential(
nn.Linear(256 * 3 * 10, 1024),
nn.Dropout(),
nn.ReLU(inplace=True),
nn.Linear(1024, 4 * 36),
)
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = out.view(out.size(0), -1) # 保留第一维度
out = self.layer5(out)
return out
class MyData(Dataset):
def __init__(self, if_train = True):
super(MyData, self).__init__()
self.transforms = transforms.Compose(
[
transforms.Resize((60, 160)),
transforms.Grayscale(),
transforms.ToTensor()
]
)
if if_train == True:
self.root = "./dataset5/train/"
else:
self.root = "./dataset5/test/"
def __len__(self):
return(len(os.listdir(self.root)))
def __getitem__(self, item):
img_path = self.root + os.listdir(self.root)[item]
img = Image.open(img_path)
img = self.transforms(img)
label = os.listdir(self.root)[item].split('_')[0]
label = text2vec(label)
label = label.view(1, -1)[0] # 极其重要
# t = torch.zeros(100, 1, 4, 36)
# print(t.view(1, -1))
# print(t.view(1, -1)[0])
return img, label # img torch.Size([1, 60, 160])
if __name__ == '__main__':
# 开始训练
m = MyNet()
# 定义 loss 函数
loss_fn = nn.MultiLabelSoftMarginLoss()
optimizer = torch.optim.Adam(m.parameters(), lr=0.0008)
total_step = 0
count = 80000
count_a = 0
for epoch in range(80000):
train_data = MyData()
test_data = MyData(if_train=False)
# 使用 pytorch 自带的 DataLoader 定义一个数据迭代器
train_dataloader = DataLoader(train_data, batch_size=10, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=10, shuffle=False)
for imgs, targets in train_dataloader:
imgs = Variable(imgs)
count_a += 1
if count_a < count:
m.train()
outputs = m(imgs)
loss = loss_fn(outputs, targets)
outputs = outputs.view(-1, 4, 36)
# print('预测为', batch_vec_2_text(outputs))
targets = targets.view(-1, 4, 36)
# print('实际为', batch_vec_2_text(targets))
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_step += 1
print("训练{}次,loss:{}".format(total_step * 1, loss.item()),compare_list(batch_vec_2_text(outputs), batch_vec_2_text(targets)))
with open('log.txt', 'a+') as f:
f.write("训练{}次,loss:{}".format(total_step * 1, loss.item()) + str(compare_list(batch_vec_2_text(outputs), batch_vec_2_text(targets))) + "\n")
f.close()
else:
break
imgs, targets = next(iter(test_dataloader))
imgs = Variable(imgs)
m.eval()
outputs = m(imgs)
loss = loss_fn(outputs, targets)
outputs = outputs.view(-1, 4, 36)
# print('样本预测为', batch_vec_2_text(outputs))
targets = targets.view(-1, 4, 36)
# print('样本实际为', batch_vec_2_text(targets))
print("训练{}次,样本loss:{}".format(total_step * 1, loss.item()),
compare_list(batch_vec_2_text(outputs), batch_vec_2_text(targets)))
with open('log1.txt', 'a+') as f:
f.write("训练{}次,loss:{}".format(total_step * 1, loss.item()) + str(
compare_list(batch_vec_2_text(outputs), batch_vec_2_text(targets))) + "\n")
f.close()
if count_a%15000 == 0:
torch.save(m, "model.pth")
torch.save(m, "model.pth")
