update: README.md

README.md

@@ -31,6 +31,10 @@ pip install -r requirements.txt
 
 3. 根据提示生成数据集，生成3次数据集，分别用于训练，用于测试，用于验证。
 
+```shell
+python captcha_gen.py
+```
+
 建议的数据集长度如下：
 
 | 数据集    | 长度     |
@@ -39,10 +43,6 @@ pip install -r requirements.txt
 | Test     | 1000    |
 | Predict  | 30      |
 
-```shell
-python captcha_gen.py
-```
-
 4. 训练模型
 
 ```shell

predict.py

@@ -29,8 +29,9 @@ def main():
     model.eval()
 
     # random pickup some test images
-    pickup_count = 30
-    pickup_rect = [5, 6]
+    pickup_count = int(input("输入验证集的图片长度（默认30）：") or 30)
+    input_rect = input("输入图片的行数和列数（默认5x6）：")
+    pickup_rect = [int(i) for i in input_rect.split("x")] if input_rect else [5, 6]
     files = os.listdir(captcha_settings.PREDICT_DATASET_PATH)
     images_picked = random.sample(files, pickup_count)
     
@@ -38,7 +39,7 @@ def main():
     fig, axes = plt.subplots(nrows=pickup_rect[0], ncols=pickup_rect[1], figsize=(10, 8))
     for i, image_name in enumerate(images_picked):
         real_text = image_name.split(".")[0].split("_")[-1]
-        file_path = os.path.join(captcha_settings.TEST_DATASET_PATH, image_name)
+        file_path = os.path.join(captcha_settings.PREDICT_DATASET_PATH, image_name)
         pred_text = predict(model, file_path)
         correct = real_text == pred_text
         axes[i//pickup_rect[1], i%pickup_rect[1]].imshow(plt.imread(file_path))