【项目04】基于卷积神经网络N的图片识别 —— 车牌识别¶

##################################################################################
# 数据集预处理
# 作者: Xinyu Ou (http://ouxinyu.cn)
# 数据集名称：车牌识别数据集
# 数据集简介: VehicleLicense车牌识别数据集包含16151张单字符数据，所有的单字符均为严格切割且都转换为黑白二值图像。
# 本程序功能:
# 1. 将数据集按照7:1:2的比例划分为训练验证集、训练集、验证集、测试集
# 2. 代码将生成4个文件：训练验证集trainval.txt, 训练集列表train.txt, 验证集列表val.txt, 测试集列表test.txt, 数据集信息dataset_info.json
# 3. 代码输出信息：图像列表已生成, 其中训练验证集样本12877，训练集样本11232个, 验证集样本1645个, 测试集样本3274个, 共计16151个。
# 4. 生成数据集标签词典时，需要根据标签-文件夹列表匹配标签列表
###################################################################################

import os
import json
import codecs

num_trainval = 0
num_train = 0
num_val = 0
num_test = 0
class_dim = 0
dataset_info = {
    'dataset_name': '',
    'num_trainval': -1,
    'num_train': -1,
    'num_val': -1,
    'num_test': -1,
    'class_dim': -1,
    'label_dict': {}
}
# 本地运行时，需要修改数据集的名称和绝对路径，注意和文件夹名称一致
dataset_name = 'VehicleLicense'
dataset_path = 'D:\\Workspace\\ExpDatasets\\'
dataset_root_path = os.path.join(dataset_path, dataset_name)
excluded_folder = ['.DS_Store', '.ipynb_checkpoints']      # 被排除的文件夹

# 获取标签和文件夹的对应关系，即省市拼音和中文对照关系
json_label_match = os.path.join(dataset_root_path, 'label_match.json')
label_match = json.loads(open(json_label_match, 'r', encoding='utf-8').read()) 
    
# 定义生成文件的路径
data_path = os.path.join(dataset_root_path, 'Data')
trainval_list = os.path.join(dataset_root_path, 'trainval.txt')
train_list = os.path.join(dataset_root_path, 'train.txt')
val_list = os.path.join(dataset_root_path, 'val.txt')
test_list = os.path.join(dataset_root_path, 'test.txt')
dataset_info_list = os.path.join(dataset_root_path, 'dataset_info.json')

# 检测数据集列表是否存在，如果存在则先删除。其中测试集列表是一次写入，因此可以通过'w'参数进行覆盖写入，而不用进行手动删除。
if os.path.exists(trainval_list):
    os.remove(trainval_list)
if os.path.exists(train_list):
    os.remove(train_list)
if os.path.exists(val_list):
    os.remove(val_list)
if os.path.exists(test_list):
    os.remove(test_list)

# 按照比例进行数据分割
class_name_list = os.listdir(data_path)
with codecs.open(trainval_list, 'a', 'utf-8') as f_trainval:
    with codecs.open(train_list, 'a', 'utf-8') as f_train:
        with codecs.open(val_list, 'a', 'utf-8') as f_val:
            with codecs.open(test_list, 'a', 'utf-8') as f_test:
                for class_name in class_name_list:
                    if class_name not in excluded_folder:
                        dataset_info['label_dict'][str(class_dim)] = label_match[class_name]        # 按照文件夹名称和label_match进行标签匹配
                        images = os.listdir(os.path.join(data_path, class_name))
                        count = 0
                        for image in images:
                            if count % 10 == 0:  # 抽取大约10%的样本作为验证数据
                                f_val.write("{0}\t{1}\n".format(os.path.join(data_path, class_name, image), class_dim))
                                f_trainval.write("{0}\t{1}\n".format(os.path.join(data_path, class_name, image), class_dim))
                                num_val += 1
                                num_trainval += 1
                            elif count % 10 == 1 or count % 10 == 2:  # 抽取大约20%的样本作为测试数据
                                f_test.write("{0}\t{1}\n".format(os.path.join(data_path, class_name, image), class_dim))
                                num_test += 1
                            else:
                                f_train.write("{0}\t{1}\n".format(os.path.join(data_path, class_name, image), class_dim))
                                f_trainval.write("{0}\t{1}\n".format(os.path.join(data_path, class_name, image), class_dim))
                                num_train += 1
                                num_trainval += 1
                            count += 1
                    class_dim += 1        

# 将数据集信息保存到json文件中供训练时使用
dataset_info['dataset_name'] = dataset_name
dataset_info['num_trainval'] = num_trainval
dataset_info['num_train'] = num_train
dataset_info['num_val'] = num_val
dataset_info['num_test'] = num_test
dataset_info['class_dim'] = class_dim

with codecs.open(dataset_info_list, 'w', encoding='utf-8') as f_dataset_info:
    json.dump(dataset_info, f_dataset_info, ensure_ascii=False, indent=4, separators=(',', ':')) # 格式化字典格式的参数列表

print("图像列表已生成, 其中训练验证集样本{}，训练集样本{}个, 验证集样本{}个, 测试集样本{}个, 共计{}个。".format(num_trainval, num_train, num_val, num_test, num_train+num_val+num_test))

# 定义数据集映射函数获取数据的图像矩阵和label
def data_mapper(sample):
    img, label = sample
    img = cv2.imread(img, 0)
    img = cv2.resize(img, (img_size, img_size))       # 将图像尺度resize为指定尺寸
    img = np.array(img).reshape(1, img_channel, img_size, img_size).astype('float32')   # 将图像数据类型转化为float32
#     img = img.transpose((2, 0, 1))          # 调整数据形状paddle默认格式(通道,高度,宽度)
    img = img/255.0                         # 将像素值归一化到[0,1]之间
    
    return img, label

# 定义数据集reader，用于从列表文件中批量获取图像
def data_reader(data_list_path):
    #定义读取函数，从列表文件中读取
    def reader():
        with open(data_list_path, 'r') as f:
            lines = f.readlines()
            for line in lines:
                img_path, label = line.split('\t')
                yield img_path, int(label)
    #使用多线程方式，通过用户自定义的映射器mapper来映射reader返回的样本（到输出队列） 
    return paddle.reader.xmap_readers(data_mapper, reader, cpu_count(), 512) 

# 测试：输出第0个batch的数据形态
for batch_id, data in enumerate(train_reader()):
    print(data[0][0].shape)
    break

# 绘制训练batch精度和平均loss
def draw_process(title, loss_label, accuracy_label, iters, losses, accuracies, figure_path=None, figurename=None, isShow=False):

    # 1.第一组坐标轴 Loss
    _, ax1 = plt.subplots()   # plt.subplots(figsize=(14,6))
    ax1.plot(iters, losses, color='red', label=loss_label)
    ax1.set_xlabel('Iters', fontsize=20)
    ax1.set_ylabel(loss_label, fontsize=20)
    max_loss = max(losses)
    ax1.set_ylim(0, max_loss*1.2)
    
    # 2.第二组坐标轴 accuracy
    ax2 = ax1.twinx()
    ax2.plot(iters, accuracies, color='blue', label=accuracy_label)
    ax2.set_ylabel(accuracy_label, fontsize=20)
    max_acc = max(accuracies)
    ax2.set_ylim(0, max_acc*1.2)
    
    # 3.配置图例
    plt.title(title, fontsize=24)
    handles1, labels1 = ax1.get_legend_handles_labels()
    handles2, labels2 = ax2.get_legend_handles_labels()
    plt.legend(handles1+handles2, labels1+labels2, loc='best')
    plt.grid()
    
    # 4.将绘图结果保存到 final_figures 目录
    plt.savefig(os.path.join(figure_path, figurename + '.png'))
    
    # 5.显示绘图结果
    if isShow is True:
        plt.show()
        
### 测试可视化函数 ###################################################
if __name__ == '__main__':
    try:
        train_log = json.loads(open(os.path.join(final_figures_path, 'train.json'), 'r', encoding='utf-8').read())
        draw_process('Training', 'loss', 'accuracy', train_log['iters'], train_log['losses'], train_log['accs_top1'], figure_path=final_figures_path, figurename='train', isShow=True)
    except:
        print('数据不存在，无法进行绘制')        

# 定义多层感知机(CNN)
class CNN(fluid.dygraph.Layer):
    name_scope = 'CNN'
    def __init__(self, num_classes=65): # 初始化CNN类,并为CNN增加对象self.x
        super(CNN, self).__init__()
        self.conv1 = Conv2D(num_channels=1, num_filters=28, filter_size=5, stride=1, act='relu')
        self.pool1 = Pool2D(pool_size=2, pool_stride=1, pool_type='max')
        self.conv2 = Conv2D(num_channels=28, num_filters=32, filter_size=3, stride=1, act='relu')
        self.pool2 = Pool2D(pool_size=2, pool_stride=1, pool_type='max')
        self.conv3 = Conv2D(num_channels=32, num_filters=32, filter_size=3, stride=1, act='relu')
        self.fc1 = Linear(input_dim=32*10*10, output_dim=num_classes)
 
    def forward(self,input): # 为CN类增加forward方法
#         print(input.shape)
        x = self.conv1(input)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.pool2(x)
        x = self.conv3(x)
        x = fluid.layers.reshape(x, [x.shape[0], -1])  # x_shape=[32*10*10, -1]
        y = self.fc1(x)

        return y
    
# Best result: epoch=94, 
# 测试集精度accuracy=[0.97664]
if __name__ == '__main__':
    model = CNN()
    paddle.summary(model, (10,1,20,20))

if __name__ == '__main__':
    # 初始化绘图列表
    train_log = {'iters': [], 'losses': [], 'accs_top1': []}
    val_log = {'iters': [], 'losses': [], 'accs_top1': []}
    
    with fluid.dygraph.guard(PLACE): 
        model = CNN() #模型实例化
        
        # 启动训练过程
        train(model) 
        
# 训练完成，最终性能accuracy=0.97864(epoch=10), 总耗时52.69s, 已将其保存为：best_model
# 训练完成，最终性能accuracy=0.98348(epoch=19), 总耗时105.66s, 已将其保存为：best_model

with fluid.dygraph.guard(PLACE):
    model_dict, _ = fluid.load_dygraph(os.path.join(final_models_path, 'best_model'))
    model = CNN() #模型实例化
    model.load_dict(model_dict) #加载模型参数    
 
    #启动训练过程
    _, avg_acc = test(model, test_reader())            
    print('测试集精度为:{:.5f}'.format(avg_acc))    

def load_image(img_path):
    img = cv2.imread(img_path, 0)                                                        # cv2.imread(path, 0|1)，其中0表示灰度模式，1表示彩色模式
    img = cv2.resize(img, (img_size, img_size))                                          # resize image with high-quality 图像大小为28*28
    img = np.array(img).reshape(img_channel, img_size, img_size).astype('float32')    # 返回新形状的数组,把它变成一个 numpy 数组以匹配数据馈送格式。
    img = img/255.0                                                                      # 将数据归一化到[0~1]之间,也可以归一化为[-1,1]之间, img = img/255.0*2.0-1.0                       
    return img

# 将图像转换为二值模式
def color2bin(img_path):
    img_gray = cv2.imread(img_path, 0)                                                   # cv2.imread(path, 0|1)，其中0表示灰度模式，1表示彩色模式
    ret, img_bin = cv2.threshold(img_gray, 120, 255, cv2.THRESH_BINARY)                  # 将图像转换为二值模式，分割阈值为120
    return img_bin


# 定义车牌字符分割函数，实现将车牌分割成单字符
def Segmentation(img_path, img_name):
    
    img_bin = color2bin(os.path.join(img_path, img_name))
    print(img_bin.shape)

    # 对车牌图片进行处理，分割出车牌中的每一个字符并保存
    result = []
    for col in range(img_bin.shape[1]):
        result.append(0)
        for row in range(img_bin.shape[0]):
            result[col] = result[col] + img_bin[row][col]/255
    character_dict = {}
    num = 0
    i = 0
    while i < len(result):
        if result[i] == 0:
            i += 1        
        else:
            index = i + 1
            while result[index] != 0:
                index += 1
            character_dict[num] = [i, index-1]
            num += 1
            i = index
#     print(character_dict)

    for i in range(8):
        if i==2:
            continue
        padding = (170 - (character_dict[i][1] - character_dict[i][0])) / 2
        ndarray = np.pad(img_bin[:,character_dict[i][0]:character_dict[i][1]], ((0,0), (int(padding), int(padding))), 'constant', constant_values=(0,0))
        ndarray = cv2.resize(ndarray, (20,20))
        
        tmp_path = os.path.join(img_path, 'tmp')        
        if not os.path.exists(tmp_path):
            os.makedirs(tmp_path)        
        cv2.imwrite(os.path.join(tmp_path, str(i) + '.png'), ndarray)
        
######################################################################        
# 输出二值化后的图像示例
if __name__ == "__main__":
    img_name = 'test02.png'                            
    img_path = os.path.join(dataset_root_path, 'Infer')

    img_bin = color2bin(os.path.join(img_path, img_name))
    print(img_bin)
    plt.imshow(img_bin, cmap='gray')

# 0. 设置待预测样本
img_name = 'test02.png'                            
img_path = os.path.join(dataset_root_path, 'Infer') 

# 1. 获取标签名称和标签ID的对应关系
json_label_match = os.path.join(dataset_root_path, 'dataset_info.json')
label_match = json.loads(open(json_label_match, 'r', encoding='utf-8').read()) 

# 2. 构建预测动态图过程
with fluid.dygraph.guard():
    model = CNN()#模型实例化
    model_dict,_= fluid.load_dygraph(os.path.join(final_model_path, 'best_model'))
    model.load_dict(model_dict)#加载模型参数
    model.eval()#评估模式
    
    Segmentation(img_path, img_name) # 原始车牌图片切割成单字符，并进行二值化处理
    lab = [] # 定义每个预测车牌的预测标签
    for i in range(8):
        if i==2:
            continue
        infer_imgs = []
        infer_imgs.append(load_image(os.path.join(img_path, 'tmp', str(i) + '.png')))
        infer_imgs = np.array(infer_imgs)
        infer_imgs = fluid.dygraph.to_variable(infer_imgs)
        result = model(infer_imgs)
        lab.append(np.argmax(result.numpy()))

# 3. 输出预测结果
print('\n车牌识别结果为：',end='')
for i in range(len(lab)):
    print(label_match['label_dict'][str(lab[i])], end='')
display(Image.open(os.path.join(img_path, img_name)))