验证码识别与之前的几个任务不同，这是一个多标签的分类的任务，也就是是一个数据对应着几个标签，只有所用的标签都预测对时，才算真正的预测成功了。
一. 数据的准备工作
与以往不同，这次的数据，我们是利用python的第三方库来生成验证码图片，下面进行代码演示，非常简单。
a = ['1','2','3','4'] img = ImageCaptcha() captcha=img.generate(a) #生成图片，根据a中的内容 captcha_image = PIL.Image.open(captcha) #读取图片 captcha_image.show() #显示图片
网络异常，图片无法展示
|
2. 基于上面的代码我们可以很轻松装备几万张数据，然后用一个Excel保存图片以及图片对应的标签，这里就不作代码展示了，展示一下Excel文件。
3. 下面开始读取我们的Excel文件，来构造我们的数据集，标签有4个，这里我们需要one-hot编码一下，弄成长度为40的向量，这也是一个需要特别需要注意的地方。
. 读取csv文件
def read_data(): data = pd.read_csv("qwe.csv") img_path = data["ID"].values label = data.iloc[:,data.columns!="ID"].values y = [] for x in label: t = one_hot(x) y.append(np.array(t)) return img_path,np.array(y)
进行one-hot编码
def one_hot(x): tmp = [0 for i in range(40)] for step,i in enumerate(x): tmp[i+10*step] = 1 return tmp
最后构造DataLoader，与显示最后的标签形式，到这里数据的准备工作就基本上完成了。
class DataSet(Dataset): def __init__(self): self.img_path,self.label = read_data() def __getitem__(self, index): img_path = self.img_path[index] img = cv2.imread(img_path,0) img = img/255. img = torch.from_numpy(img).float() img = torch.unsqueeze(img,0) label = torch.from_numpy(self.label[index]).float() return img,label def __len__(self): return len(self.img_path) data = DataSet() data_loader = DataLoader(data,shuffle=True,batch_size=64,drop_last=True)
二. 网络的构建与优化、损失函数的选取以及训练
1.网络的构建和优化函数在这里就不做多的说明了，直接看代码
class CNN_Network(nn.Module): def __init__(self): super(CNN_Network, self).__init__() self.layer1 = nn.Sequential( nn.Conv2d(1, 16, stride=1, kernel_size=3, padding=1), nn.BatchNorm2d(16), nn.ReLU(inplace=True) ) self.layer2 = nn.Sequential( nn.Conv2d(16, 32, stride=1, kernel_size=3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.MaxPool2d(stride=2, kernel_size=2), # 30 80 ) self.layer3 = nn.Sequential( nn.Conv2d(32, 64, stride=1, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.Conv2d(64,128,kernel_size=3,stride=1,padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2,stride=2), # 15 40 ) self.fc = nn.Sequential( nn.Linear(128 * 15 * 40, 2048), nn.ReLU(inplace=True), nn.Linear(2048, 1024), nn.ReLU(inplace=True), nn.Linear(1024, 40) ) def forward(self, x): x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = x.view(x.size(0), -1) x = self.fc(x) return x model = CNN_Network() optimizer = torch.optim.Adam(model.parameters(),lr=0.001)
2.损失函数我们这里使用的是多标签分类的损失函数，和交叉熵损失函数的公式比较相像，
l o s s ( x , y ) = − ∑ i y [ i ] ∗ l o g ( 11 + e x p ( − x [ i ] ) + ( 1 − y [ i ] ) ∗ l o g ( e x p ( − x [ i ] ) 1 + e x p ( − x [ i ] ) )
,上面是公式，有兴趣的可以自己研究一下。
error = nn.MultiLabelSoftMarginLoss() #注意输入的数据要是float32类型的，否则会出错。
3.所有的准备工作都完成了，下面就开始训练吧。
for i in range(2): for x_index,y in data_loader: pass x = Variable(x_index) optimizer.zero_grad() label = Variable(y) out = model(x) loss = error(out,label) print(loss) loss.backward() optimizer.step() torch.save(model.state_dict(),"验证码识别.pth")
三. 测试模型
1.训练完成后，来测试一下我们的模型吧.
cnn = CNN_Network() cnn.load_state_dict(torch.load("验证码识别.pth")) a = cv2.imread("./data/9354.jpg",0) b = cv2.resize(a,(200,200)) cv2.imshow('a',b) cv2.waitKey(0) a = a/255. a = torch.from_numpy(a).float() a = torch.unsqueeze(a,0) a = torch.unsqueeze(a,0) pred = cnn(a) print(pred.size()) a1 = torch.argmax(pred[0,:10],dim=0) #第一个标签 a2 = torch.argmax(pred[0,10:20],dim=0) #第二个标签 a3 = torch.argmax(pred[0,20:30],dim=0) #第三个标签 a4 = torch.argmax(pred[0,30:],dim=0) #第四的标签 pred = [a1,a2,a3,a4] print(pred)
预测的图片
预测结果
上面就完成了所有的工作了。
Thank for your reading !!!
公众号：FPGA之旅
