基于多任务学习模型用于验证码的分类预测功能实现

import tensorflow as tf from tensorflow.keras.models import load_model import matplotlib.pyplot as plt import os import numpy as np import string # 载入之前训练好的模型 model = load_model('Best_Captcha_tfdata.h5') # 字符包含所有数字和所有大小写英文字母，一共 62 个 characters = string.digits + string.ascii_letters # 类别数 num_classes = len(characters) # 批次大小 batch_size = 64 # 测试集数据 test_dir = "./captcha/test/" # 获取所有验证码图片路径和标签 def get_filenames_and_classes(dataset_dir): # 存放图片路径 photo_filenames = [] # 存放图片标签 y = [] for filename in os.listdir(dataset_dir): # 获取文件完整路径 path = os.path.join(dataset_dir, filename) # 保存图片路径 photo_filenames.append(path) # 取文件名前 4 位，也就是验证码的标签 captcha_text = filename[0:4] label = np.zeros((4, num_classes), dtype=np.uint8) # one-hot编码 for i, ch in enumerate(captcha_text): label[i, characters.find(ch)] = 1 y.append(label) # 返回该路径下图片路径和标签 return np.array(photo_filenames), np.array(y) # 获取测试集图片路径和标签 x_test,y_test = get_filenames_and_classes(test_dir) # 图像处理函数 def image_function(filenames, label): # 根据图片路径读取图片内容 image = tf.io.read_file(filenames) # 将图像解码为 jpeg 格式的 3 维数据 image = tf.image.decode_jpeg(image, channels=3) # 归一化 image = tf.cast(image, tf.float32) / 255.0 # 返回图片数据和标签 return image, label # 标签处理函数 # 获得每一个批次的图片数据和标签 def label_function(image, label): # 原来的shape: (64,4,62) 进行tf.transpose后，得到(4,64,62) label = tf.transpose(label, [1, 0, 2]) # image: (64,x,y,3) label[0] (64,62) return image, (label[0], label[1], label[2], label[3]) # 创建 dataset 对象，传入测试集图片路径和标签 dataset_test = tf.data.Dataset.from_tensor_slices((x_test, y_test)) dataset_test = dataset_test.shuffle(buffer_size=1000,reshuffle_each_iteration=True) dataset_test = dataset_test.map(image_function) dataset_test = dataset_test.repeat(1) dataset_test = dataset_test.batch(batch_size) # (x,y) in dataset_test: x:(batch_size, , , ) y:(4,batch_size,num_calss) dataset_test = dataset_test.map(label_function) # 用于统计准确率 acc_sum = 0 # 统计批次数量 n = 0 /* 开发不易，整理也不易，如需要详细的说明文档和程序，以及完整的数据集，训练好的模型，或者进一步开发， 可加作者新联系方式咨询，WX：Q3101759565，QQ：3101759565 */ # 获得一个批次数据 x,y = next(iter(dataset_test)) # 预测结果 # pred:[4,batch_size,numclass] pred = model.predict(x) # 获得对应编号 # pred:[4,batch_size] pred = np.argmax(pred, axis=-1) # (4,64)->(64,4) pred = pred.T # 获得标签数据 # y： [4,batch_size,numclass] # label: [4,batch_size] label = np.argmax(y, axis=-1) # (4,64)->(64,4) label = label.T # (64,) pred = decode_batch(pred) # (64,) label = decode_batch(label) # 获取前 3 张图片数据 for i,image in enumerate(x[:3]): plt.imshow(image) plt.title('real:%s\npred:%s' % (label[i], pred[i])) plt.show()