update readme

README.md

@@ -28,6 +28,7 @@
 验证码样例
 =====
 ![](https://raw.githubusercontent.com/dee1024/pytorch-captcha-recognition/master/docs/number.png)
+![](https://raw.githubusercontent.com/dee1024/pytorch-captcha-recognition/master/docs/number2.png)
 
 
 快速开始

captcha-gen.py → captcha_gen.py

@@ -3,18 +3,12 @@
 from PIL import Image
 import random
 import time
-import string
+import captcha_setting
 
-# 验证码中的字符
-# string.digits + string.ascii_uppercase
-NUMBER = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
-ALPHABET = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z']
-char_set = NUMBER
-
-def random_captcha(char_set=char_set, captcha_size=4):
+def random_captcha():
     captcha_text = []
-    for i in range(captcha_size):
-        c = random.choice(char_set)
+    for i in range(captcha_setting.MAX_CAPTCHA):
+        c = random.choice(captcha_setting.ALL_CHAR_SET)
         captcha_text.append(c)
     return ''.join(captcha_text)
 
@@ -27,7 +21,7 @@ def gen_captcha_text_and_image():
 
 if __name__ == '__main__':
     count = 5
-    path = 'dataset/predict/1'
+    path = captcha_setting.PREDICT_DATASET_PATH+'/1'
     for i in range(count):
         now = str(int(time.time()))
         text, image = gen_captcha_text_and_image()

captcha-predict.py → captcha_predict.py

@@ -6,13 +6,7 @@
 from torch.autograd import Variable
 import torchvision.datasets as dsets
 from visdom import Visdom # pip install Visdom
-
-vis = Visdom()
-NUMBER = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
-ALPHABET = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z']
-char_set = NUMBER
-MAX_CAPTCHA = 4
-CHAR_SET_LEN = len(char_set)
+import captcha_setting
 
 class CNN(nn.Module):
     def __init__(self):
@@ -40,7 +34,7 @@ def __init__(self):
             nn.Dropout(0.5),  # drop 50% of the neuron
             nn.ReLU())
         self.rfc = nn.Sequential(
-            nn.Linear(1024, MAX_CAPTCHA*CHAR_SET_LEN),
+            nn.Linear(1024, captcha_setting.MAX_CAPTCHA * captcha_setting.ALL_CHAR_SET_LEN),
         )
 
     def forward(self, x):
@@ -64,25 +58,20 @@ def main():
 
     pre_loader = torch.utils.data.DataLoader(dataset=pre_dataset,batch_size=1,shuffle=False)
 
-
-    # print(char_set)
-
+    vis = Visdom()
     for images in pre_loader:
         image = images[0]
         # print(image.shape)
         vimage = Variable(image)
         predict_label = cnn(vimage)
-        # c0 = char_set[np.argmax(predict_label[0, 0:36].data.numpy())]
-        # c1 = char_set[np.argmax(predict_label[0, 36:72].data.numpy())]
-        # c2 = char_set[np.argmax(predict_label[0, 72:108].data.numpy())]
-        # c3 = char_set[np.argmax(predict_label[0, 108:144].data.numpy())]
 
-        c0 = char_set[np.argmax(predict_label[0, 0:10].data.numpy())]
-        c1 = char_set[np.argmax(predict_label[0, 10:20].data.numpy())]
-        c2 = char_set[np.argmax(predict_label[0, 20:30].data.numpy())]
-        c3 = char_set[np.argmax(predict_label[0, 30:40].data.numpy())]
+        c0 = captcha_setting.ALL_CHAR_SET[np.argmax(predict_label[0, 0:captcha_setting.ALL_CHAR_SET_LEN].data.numpy())]
+        c1 = captcha_setting.ALL_CHAR_SET[np.argmax(predict_label[0, captcha_setting.ALL_CHAR_SET_LEN:2 * captcha_setting.ALL_CHAR_SET_LEN].data.numpy())]
+        c2 = captcha_setting.ALL_CHAR_SET[np.argmax(predict_label[0, 2 * captcha_setting.ALL_CHAR_SET_LEN:3 * captcha_setting.ALL_CHAR_SET_LEN].data.numpy())]
+        c3 = captcha_setting.ALL_CHAR_SET[np.argmax(predict_label[0, 3 * captcha_setting.ALL_CHAR_SET_LEN:4 * captcha_setting.ALL_CHAR_SET_LEN].data.numpy())]
 
         c = '%s%s%s%s' % (c0, c1, c2, c3)
+        print(c)
         vis.images(image, opts=dict(caption=c))
 
 if __name__ == '__main__':

captcha_setting.py

@@ -0,0 +1,16 @@
+# -*- coding: UTF-8 -*-
+import string
+# 验证码中的字符
+# string.digits + string.ascii_uppercase
+NUMBER = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
+ALPHABET = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z']
+
+ALL_CHAR_SET = NUMBER + ALPHABET
+ALL_CHAR_SET_LEN = len(ALL_CHAR_SET)
+MAX_CAPTCHA = 4
+
+
+
+TRAIN_DATASET_PATH = 'dataset/train'
+TEST_DATASET_PATH = 'dataset/test'
+PREDICT_DATASET_PATH = 'dataset/predict'

mydataset.py

@@ -5,9 +5,6 @@
 from PIL import Image
 import onehotencoding as ohe
 
-MAX_CAPTCHA = 4     # 验证码数字的个数
-CHAR_SET_LEN = 36   # 每个验证码的可能出现的的值：0，1，2，3，4，5，6，7，8，9，一共 10个
-
 class mydataset(Dataset):
 
     def __init__(self, folder, transform=None):
@@ -29,6 +26,6 @@ def __getitem__(self, idx):
 def get_data_loader():
     normalize = transforms.Normalize(mean=[0.92206, 0.92206, 0.92206], std=[0.08426, 0.08426, 0.08426])
     transform = transforms.Compose([transforms.ToTensor(),normalize])
-    dataset =  mydataset(transform=transform)
-    return DataLoader(dataset,batch_size=32, shuffle=True)
+    dataset = mydataset(transform=transform)
+    return DataLoader(dataset,batch_size=64, shuffle=True)
 

onehotencoding.py

@@ -1,11 +1,10 @@
 # -*- coding: UTF-8 -*-
 import numpy as np
+import captcha_setting
 
-len = 36
-count = 4
 
 def encode(text):
-    vector = np.zeros(count*len, dtype=float)
+    vector = np.zeros(captcha_gen.CHAR_SET_LEN*captcha_gen.MAX_CAPTCHA, dtype=float)
     def char2pos(c):
         if c =='_':
             k = 62
@@ -19,7 +18,7 @@ def char2pos(c):
                     raise ValueError('error')
         return k
     for i, c in enumerate(text):
-        idx = i * len + char2pos(c)
+        idx = i * captcha_setting.ALL_CHAR_SET_LEN + char2pos(c)
         vector[idx] = 1.0
     return vector
 
@@ -28,7 +27,7 @@ def decode(vec):
     text=[]
     for i, c in enumerate(char_pos):
         char_at_pos = i #c/63
-        char_idx = c % len
+        char_idx = c % captcha_setting.ALL_CHAR_SET_LEN
         if char_idx < 10:
             char_code = char_idx + ord('0')
         elif char_idx <36: