import os
import json
import codecs
import time
import paddle
from pprint import pprint

starting_time = time.strftime("%Y%m%d%H%M", time.localtime())

# 1. 定义训练参数
train_parameters = {
    'dataset_name': 'Butterfly',       # 数据集名称
    'architecture': 'ResNet50_pretrained',        # 模型名称
    'training_data': 'trainval',             # 训练所用数据 train|trainval
    'starting_time': starting_time,          # 全局启动时间
    'input_size': [3, 224, 224],             # 输入样本的尺度
    'mean_value': [0.485, 0.456, 0.406],     # Imagenet均值
    'std_value': [0.229, 0.224, 0.225],      # Imagenet标准差    
    'num_trainval': -1,  
    'num_train': -1,
    'num_val': -1,
    'num_test': -1,
    'class_dim': -1,
    'label_dict': {},
    'total_epoch': 2,                # 总迭代次数, 代码调试好后考虑
    'batch_size': 64,                 # 设置每个批次的数据大小，同时对训练提供器和测试
    'log_interval': 1,                # 设置训练过程中，每隔多少个batch显示一次
    'eval_interval': 1,               # 设置每个多少个epoch测试一次
    'checkpointed': False,            # 是否保存checkpoint模型
    'checkpoint_train': False,          # 是否接着上一次保存的参数接着训练，优先级高于预训练模型
    'checkpoint_model':'Butterfly_ResNet50_pretrained_final',            # 设置恢复训练时载入的模型参数
    'checkpoint_time': '202102182058',   #  恢复训练时所指向的指定时间戳
    'pretrained': True,              # 是否使用预训练的模型    
    'pretrained_model':'API',              # 设置预训练模型, API|Butterflies_AlexNet_final
    'dataset_root_path': 'D:\\Workspace\\ExpDatasets',
    'result_root_path': 'D:\\Workspace\\ExpResults',
    'project_result_path': 'D:\\Workspace\\MyProjects\\ButterflyClassification\\results',
    'useGPU': True,                   # True | Flase
    'learning_strategy': {            # 学习率和优化器相关参数
        'optimizer_strategy': 'Momentum',                   # 优化器：Momentum, RMS, SGD, Adam
        'learning_rate_strategy': 'CosineAnnealingDecay',   # 学习率策略: 固定fixed, 分段衰减PiecewiseDecay, 余弦退火CosineAnnealingDecay, 指数ExponentialDecay, 多项式PolynomialDecay
        'learning_rate': 0.01,                              # 固定学习率
        'momentum': 0.9,                                    # 动量
        'Piecewise_boundaries': [60, 80, 90],               # 分段衰减：变换边界，每当运行到epoch时调整一次
        'Piecewise_values': [0.01, 0.001, 0.0001, 0.00001], # 分段衰减：步进学习率，每次调节的具体值
        'Exponential_gamma': 0.9,                           # 指数衰减：衰减指数
        'Polynomial_decay_steps': 10,                       # 多项式衰减：衰减周期，每个多少个epoch衰减一次
        'verbose': False
    },
    'augmentation_strategy': { 
        'withAugmentation': True,            # 数据扩展相关参数
        'augmentation_prob': 0.5,            # 设置数据增广的概率
        'rotate_angle': 15,                  # 随机旋转的角度
        'Hflip_prob': 0.5,                   # 随机翻转的概率
        'brightness': 0.4, 
        'contrast': 0.4,
        'saturation': 0.4,
        'hue': 0.4,
    },
}

args = train_parameters

# 2. 定义组合参数
model_name = args['dataset_name'] + '_' + args['architecture']

# 3. 定义设备选择函数: [GPU|CPU]
def init_device(useGPU=args['useGPU']):
    paddle.set_device('gpu:0') if useGPU else paddle.set_device('cpu')
init_device()

# 4. 定义数据集路径
dataset_root_path = os.path.join(args['dataset_root_path'], args['dataset_name'])
dataset_info_list = os.path.join(dataset_root_path, 'dataset_info.json')

# 5. 定义结果路径：模型、训练、日志结果图
# 定义训练时模型的路径设置
result_root_path = os.path.join(args['result_root_path'], model_name + '_' + args['starting_time'])
final_models_path = os.path.join(result_root_path, 'final_models')     # 最终用于部署和推理的模型
final_figures_path = os.path.join(result_root_path, 'final_figures')    # 训练过程曲线图
logs_path = os.path.join(result_root_path, 'logs')             # 训练过程日志
# checkpoint_ 路径用于定义恢复训练所用的模型保存
checkpoint_path = os.path.join(result_root_path, model_name + '_final')
checkpoint_load_model = os.path.join(args['result_root_path'], model_name + '_' + args['checkpoint_time'], 'checkpoint_models', args['checkpoint_model'])

# project_ 路径用于定义验证和预测时调用模型的路径
project_result_path = os.path.join(args['project_result_path'], model_name)
project_checkpoint_path = os.path.join(project_result_path, 'checkpoint_models', model_name + '_final')
project_final_model_path = os.path.join(project_result_path, 'final_models', model_name + '_final')
project_final_figure_path = os.path.join(project_result_path, 'final_figures')
project_pretrained_model = os.path.join(project_result_path, 'pretrained_dir', args['pretrained_model'])

# 初始化结果目录
def init_result_path():
    if not os.path.exists(final_models_path):
        os.makedirs(final_models_path)
    if not os.path.exists(final_figures_path):
        os.makedirs(final_figures_path)
    if not os.path.exists(logs_path):
        os.makedirs(logs_path)
    if not os.path.exists(checkpoint_path):
        os.makedirs(checkpoint_path)
init_result_path()


# 5. 初始化参数
def init_train_parameters():
    dataset_info = json.loads(open(dataset_info_list, 'r', encoding='utf-8').read())    
    args['num_trainval'] = dataset_info['num_trainval']
    args['num_train'] = dataset_info['num_train']
    args['num_val'] = dataset_info['num_val']
    args['num_test'] = dataset_info['num_test']
    args['class_dim'] = dataset_info['class_dim']
    args['label_dict'] = dataset_info['label_dict']
init_train_parameters()

if __name__ == '__main__':
    pprint(train_parameters)

# 载入项目文件夹
import sys
sys.path.append(r"D:\\Workspace\\MyProjects\\ButtleflyClassification")
#################################################################################

import sys
import os
sys.path.append(os.path.abspath(os.path.join(sys.path[0], '..')))
import paddle
import paddle.optimizer as optimizer

train_parameters = { 
    'num_train': 200,
    'total_epoch': 20,                # 总迭代次数, 代码调试好后考虑
    'batch_size': 64, 
    'learning_strategy': {                                 # 学习率和优化器相关参数
        'optimizer_strategy': 'Momentum',                  # 优化器：Momentum, RMS, SGD, Adam
        'learning_rate_strategy': 'CosineAnnealingDecay',  # 学习率策略: 固定fixed, 分段衰减PiecewiseDecay, 余弦退火CosineAnnealingDecay, 指数ExponentialDecay, 多项式PolynomialDecay
        'learning_rate': 0.01,                             # 固定学习率
        'momentum': 0.9,                                   # 动量
        'Piecewise_boundaries': [60, 80, 90],              # 分段衰减：变换边界，每当运行到epoch时调整一次
        'Piecewise_values': [0.01, 0.001, 0.0001, 0.00001],# 分段衰减：步进学习率，每次调节的具体值
        'Exponential_gamma': 0.9,                          # 指数衰减：衰减指数
        'Polynomial_decay_steps': 10,                      # 多项式衰减：衰减周期，每个多少个epoch衰减一次
        'verbose': False
    },
}
args = train_parameters['learning_strategy']

def learning_rate_setting(verbose=args['verbose']):
    if args['learning_rate_strategy'] == 'PiecewiseDecay':
        lr = optimizer.lr.PiecewiseDecay(boundaries=args['Piecewise_boundaries'], values=args['Piecewise_values'], verbose=verbose)
    elif args['learning_rate_strategy'] == 'CosineAnnealingDecay':
        step_each_epoch = train_parameters['num_train'] // (train_parameters['batch_size'] * 2)
        T_max = step_each_epoch * train_parameters['total_epoch']
        lr = optimizer.lr.CosineAnnealingDecay(learning_rate=args['learning_rate'], T_max=T_max, verbose=verbose)
    elif args['learning_rate_strategy'] == 'ExponentialDecay':
        lr = optimizer.lr.ExponentialDecay(learning_rate=args['learning_rate'], gamma=args['Exponential_gamma'], verbose=verbose)
    elif args['learning_rate_strategy'] == 'PolynomialDecay':
        lr = optimizer.lr.PolynomialDecay(learning_rate=args['learning_rate'], decay_steps=args['Polynomial_decay_steps'], verbose=verbose)
    else:
        lr = args['learning_rate']
    return lr

def optimizer_setting(model, lr):
    if args['optimizer_strategy'] == 'Momentum':
        # 阶梯型的学习率适合比较大规模的训练数据
        opt = optimizer.Momentum(learning_rate=lr, momentum=args['momentum'], parameters=model.parameters())
    elif args['optimizer_strategy'] == 'RMS':
        # 阶梯型的学习率适合比较大规模的训练数据
        opt = optimizer.RMSProp(learning_rate=lr, parameters=model.parameters())
    elif args['optimizer_strategy'] == 'SGD':
        # loss下降相对较慢，但是最终效果不错，阶梯型的学习率适合比较大规模的训练数据
        opt = optimizer.SGD(learning_rate=lr, parameters=model.parameters())
    elif args['optimizer_strategy'] == 'Adam':
        # 能够比较快速的降低 loss，但是相对后期乏力
        opt = optimizer.Adam(learning_rate=lr, parameters=model.parameters())
    else:
        print('学习率设置错误, 请重新设置。')

    return opt

# 学习率输出测试
if __name__ == '__main__':
    print('当前学习率为: {}'.format(args['learning_rate_strategy']))
    linear = paddle.nn.Linear(10, 10)
    lr = learning_rate_setting(verbose=True)
    opt = optimizer_setting(linear, lr)
    if args['learning_rate_strategy'] == 'fixed':
        print('learning = {}'.format(args['learning_rate']))
    else:
        for epoch in range(20):
            for batch_id in range(10):
                x = paddle.uniform([10, 10])
                out = linear(x)
                loss = paddle.mean(out)
                loss.backward()
                opt.step()
                opt.clear_gradients()
#                 lr.step()    # If you update learning rate each step
            lr.step()        # If you update learning rate each epoch

# 载入项目文件夹
import sys
sys.path.append(r"D:\\Workspace\\MyProjects\\ButterflyClassification")
#################################################################################

import os
import time
import json
from utils.draw import draw_process
from utils.logger import logger
from utils.optimizer import learning_rate_setting, optimizer_setting
from utils.reader import train_reader, trainval_reader, val_reader
from config import *
from eval import *
import paddle
from paddle.static import InputSpec

# 初始配置变量
total_epoch = train_parameters['total_epoch']

# 初始化绘图列表
all_train_iters = []
all_train_losses = []
all_train_accs_top1 = []
all_train_accs_top5 = []
all_test_losses = []
all_test_iters = []
all_test_accs_top1 = []
all_test_accs_top5 = []

def train(model):  
    # 初始化临时变量
    num_batch = 0
    best_result = 0
    best_result_id = 0
    elapsed = 0

    # 根据config文件设置训练数据来源
    if train_parameters['training_data'] == 'trainval':
        data_reader = trainval_reader
    elif train_parameters['training_data'] == 'train':
        data_reader = train_reader

    for epoch in range(1, total_epoch+1):
        for batch_id, (image, label) in enumerate(data_reader()):
            num_batch += 1
            
            label = paddle.unsqueeze(label, axis=1)                        
            loss, acc = model.train_batch([image], [label])
       
            if num_batch % train_parameters['log_interval'] == 0: # 每10个batch显示一次日志，适合大数据集  
                avg_loss = loss[0][0]
                acc_top1 = acc[0][0]
                acc_top5 = acc[0][1]
                
                elapsed_step = time.perf_counter() - elapsed - start
                elapsed = time.perf_counter() - start
                logger.info('Epoch:{}/{}, batch:{}, train_loss:[{:.5f}], acc_top1:[{:.5f}], acc_top5:[{:.5f}]({:.2f}s)'
                            .format(epoch, total_epoch, num_batch, loss[0][0], acc[0][0], acc[0][1], elapsed_step))

                # 记录训练过程，用于可视化训练过程中的loss和accuracy
                all_train_iters.append(num_batch)
                all_train_losses.append(avg_loss)
                all_train_accs_top1.append(acc_top1)
                all_train_accs_top5.append(acc_top5)

        # 每隔一定周期进行一次测试
        if epoch % train_parameters['eval_interval'] == 0 or epoch == total_epoch:
            # 模型校验
            avg_loss, avg_acc_top1, avg_acc_top5 = eval(model, val_reader())            
            logger.info('[validation] Epoch:{}/{}, val_loss:[{:.5f}], val_top1:[{:.5f}], val_top5:[{:.5f}]'.format(epoch, total_epoch, avg_loss, avg_acc_top1, avg_acc_top5))

            # 记录测试过程，用于可视化训练过程中的loss和accuracy
            all_test_iters.append(epoch)
            all_test_losses.append(avg_loss)
            all_test_accs_top1.append(avg_acc_top1)            
            all_test_accs_top5.append(avg_acc_top5)
            

            # 将性能最好的模型保存为final模型
            if avg_acc_top1 > best_result:
                best_result = avg_acc_top1
                best_result_id = epoch
                 
                # finetune model 用于调优和恢复训练
                model.save(os.path.join(checkpoint_path, model_name + '_final')) 
                # inference model 用于部署和预测 
                model.save(os.path.join(final_models_path, model_name + '_final'), training=False)
                logger.info('已保存当前测试模型(epoch={})为最优模型:{}_final'.format(best_result_id, model_name)) 
            logger.info('最优top1测试精度:{:.5f} (epoch={})'.format(best_result, best_result_id))

            # 根据需要决定是否需要将每次测试的模型都进行保存(if needed)，保存模型需要耗费一定的运算时间和大量的存储资源
            # 建议在训练大型模型时，开启该选项，方便训练中断时能够及时恢复训练
            # 训练小型模型（训练时间短）时，可以关闭该选项，以进一步提高训练速度
            if train_parameters['checkpointed']:
                model.save(os.path.join(checkpoint_path, model_name + '_' + str(epoch)))

    logger.info('训练完成，最终性能accuracy={:.5f}(epoch={}), 总耗时{:.2f}s, 已将其保存为：{}_final'.format(best_result, best_result_id, time.perf_counter() - start, model_name))


if __name__ == '__main__':
    # 将此次训练的超参数进行保存
    data = json.dumps(train_parameters, indent=4, ensure_ascii=False, sort_keys=False, separators=(',', ':'))   # 格式化字典格式的参数列表
    logger.info(data)
    # 启动训练过程
    logger.info('训练参数保存完毕，使用{}模型, 训练{}数据, 训练集{}, 启动训练...'.format(train_parameters['architecture'],train_parameters['dataset_name'],train_parameters['training_data']))
    logger.info('当前模型目录为：{}'.format(model_name + '_' + train_parameters['starting_time']))

    # 设置输入样本的维度
    input_spec = InputSpec(shape=[None] + train_parameters['input_size'], dtype='float32', name='image')
    label_spec = InputSpec(shape=[None, 1], dtype='int64', name='label')
   
    # 载入官方标准模型，若不存在则会自动进行下载，pretrained=True|False控制是否使用Imagenet预训练参数
    network = paddle.vision.models.resnet50(num_classes=train_parameters['class_dim'], pretrained=train_parameters['pretrained'])  
    model = paddle.Model(network, input_spec, label_spec) 
    logger.info('模型参数信息：')
    logger.info(model.summary()) # 是否显示神经网络的具体信息
    
    if train_parameters['checkpoint_train'] == True:
        model.load(checkpoint_load_model)
        logger.info('载入{}中断模型和参数完毕，开始从checkpoint恢复训练'.format(train_parameters['architecture']))
        logger.info('checkpoint模型:{}'.format(checkpoint_load_model))
    else:
        if train_parameters['pretrained'] == False:
            logger.info('载入{}模型完毕，从初始状态开始训练'.format(train_parameters['architecture']))
        elif train_parameters['pretrained_model'] == 'API':
            logger.info('载入Imagenet-{}预训练模型完毕，开始微调训练(fine-tune)'.format(train_parameters['architecture']))
        else:
            model.load(project_pretrained_model)
            logger.info('载入自定义预训练{}模型完毕，开始微调训练(fine-tune)'.format(train_parameters['architecture']))
            logger.info('预训练模型:{}'.format(project_pretrained_model))
          
    # 设置学习率、优化器、损失函数和评价指标
    lr = learning_rate_setting()
    optimizer = optimizer_setting(model, lr)    
    model.prepare(optimizer,
                  paddle.nn.CrossEntropyLoss(),
                  paddle.metric.Accuracy(topk=(1,5)))   

    # 启动训练过程
    start = time.perf_counter()
    train(model)
    logger.info('训练完毕，结果路径{}.'.format(result_root_path))

    # 输出训练过程图
    logger.info('Done.')
    draw_process("Training Process", 'Train Loss', 'Train Accuracy(top1)', all_train_iters, all_train_losses, all_train_accs_top1, 'train')
    draw_process("Validation Results", 'Validation Loss', 'Validation Accuracy(top1)', all_test_iters, all_test_losses, all_test_accs_top1, 'val')

# 载入项目文件夹
import sys
sys.path.append(r"D:\\Workspace\\MyProjects\\ButterflyClassification")
#################################################################################

import numpy as np
import paddle
from config import train_parameters as args
from config import project_checkpoint_path
from utils.reader import val_reader, test_reader
import paddle.nn.functional as F
from paddle.static import InputSpec

__all__ = ['eval']

def eval(model, data_reader, verbose=0):
    if verbose == 1:
        print('开始评估...损失和精度均在上一个batch进行评估')

    accuracies_top1 = []
    accuracies_top5 = []
    losses = []
    n_total = 0
    
    for batch_id, (image, label) in enumerate(data_reader):
        n_batch = len(label)
        n_total = n_total + n_batch

        label = paddle.unsqueeze(label, axis=1)

        loss, acc = model.eval_batch([image], [label])
        losses.append(loss[0])                    
        accuracies_top1.append(acc[0][0]*n_batch)
        accuracies_top5.append(acc[0][1]*n_batch) 

        if verbose == 1:
            print('Batch:{}/{}, acc_top1:[{:.5f}], acc_top5:[{:.5f}]'.format(batch_id+1, len(data_reader), acc[0][0], acc[0][1])) 
                
    avg_loss = np.sum(losses)/n_total                 # loss 记录的是当前batch的累积值
    avg_acc_top1 = np.sum(accuracies_top1)/n_total    # metric 是当前batch的平均值
    avg_acc_top5 = np.sum(accuracies_top5)/n_total

    return avg_loss, avg_acc_top1, avg_acc_top5

##############################################################     
if __name__ == '__main__':
    # 设置输入样本的维度    
    input_spec = InputSpec(shape=[None] + args['input_size'], dtype='float32', name='image')
    label_spec = InputSpec(shape=[None, 1], dtype='int64', name='label')
    
    # 载入模型
    network = paddle.vision.models.mobilenet_v2(num_classes=args['class_dim'])  # 载入模型结构
    model = paddle.Model(network, input_spec, label_spec)                   # 模型实例化
    model.load(project_checkpoint_path)                                     # 载入调优模型的参数
    model.prepare(loss=paddle.nn.CrossEntropyLoss(),                        # 设置loss
                  metrics=paddle.metric.Accuracy(topk=(1,5)))               # 设置评价指标
    
    # 执行评估函数，并输出验证集样本的损失和精度
    avg_loss, avg_acc_top1, avg_acc_top5 = eval(model, val_reader(), verbose=1)
    print('[验证集] 损失: {:.5f}, top1精度:{:.5f}, top5精度为:{:.5f}'.format(avg_loss, avg_acc_top1, avg_acc_top5))    
    avg_loss, avg_acc_top1, avg_acc_top5 = eval(model, test_reader())
    print('[测试集] 损失: {:.5f}, top1精度:{:.5f}, top5精度为:{:.5f}'.format(avg_loss, avg_acc_top1, avg_acc_top5))    

# 载入项目文件夹
import sys
sys.path.append(r"D:\\Workspace\\MyProjects\\ButterflyClassification")
#################################################################################

#################################################
# 修改者: Xinyu Ou (http://ouxinyu.cn)
# 功能: 使用部署模型对测试集进行评估
# 基本功能：
# 1. 使用部署模型在测试集上进行批量预测，并输出预测结果
# 2. 使用部署模型在测试集上进行单样本预测，并对预测结果和真实结果进行对比
#################################################

import numpy as np
import random
import os
import cv2
import json
import matplotlib.pyplot as plt
import paddle
from config import project_final_model_path, dataset_root_path
from utils.reader import test_reader
import paddle.nn.functional as F
from utils.dataset import SimplePreprocessing

# 1. 使用部署模型在测试集上进行准确度评估
def test(model, data_reader):
    accs = []
    n_total = 0
    
    for batch_id, (image, label) in enumerate(data_reader):
        n_batch = len(label)
        n_total = n_total + n_batch
        
        # 将label扩展为规定的np矩阵
        label = paddle.unsqueeze(label, axis=1)
        
        logits = model(image)
        pred = F.softmax(logits)
        acc = paddle.metric.accuracy(pred, label)
        accs.append(acc.numpy()*n_batch)
    avg_acc = np.sum(accs)/n_total
    
    print('测试集的精确度: {:.5f}'.format(avg_acc))

# 2. 使用部署模型在测试集上进行单样本预测
def predict(model, image):    
    isTenCrop = True
    image = SimplePreprocessing(image, isTenCrop=isTenCrop)
    
    if isTenCrop:
        logits = model(image)
        pred = F.softmax(logits)
        pred = np.mean(pred.numpy(), axis=0) 
    else:
        image = paddle.unsqueeze(image, axis=0)
        logits = model(image)
        pred = F.softmax(logits)                
    pred_id = np.argmax(pred)
    
    return pred_id
    
##############################################################
if __name__ == '__main__':
    # 0. 载入模型
    model = paddle.jit.load(project_final_model_path)

    # 1. 计算测试集的准确度 
    test(model, test_reader())

    # 2. 输出单个样本测试结果
    # 2.1 获取待预测样本的标签信息
    dataset_info_list = os.path.join(dataset_root_path, 'dataset_info.json')
    with open(dataset_info_list, 'r') as f_info:
        dataset_info = json.load(f_info)
    
    # 2.2 从测试列表中随机选择一副图像
    test_list = os.path.join(dataset_root_path, 'test.txt')
    with open(test_list, 'r') as f_test:
        lines = f_test.readlines()
    line = random.choice(lines)
    img_path, label = line.split()
    img_path = os.path.join(dataset_root_path, 'Data', img_path)        
    # img_path = 'D:\\Workspace\\ExpDatasets\\Butterfly\\Data\\zebra\\zeb033.jpg'    
        
    image = cv2.imread(img_path, 1)
    
    # 2.4 给出待测样本的类别
    pred_id = predict(model, image) 
    
    # 将预测的label和ground_turth label转换为label name
    label_name_gt = dataset_info['label_dict'][str(label)]
    label_name_pred = dataset_info['label_dict'][str(pred_id)]
    print('待测样本的类别为：{}, 预测类别为：{}'.format(label_name_gt, label_name_pred))
    
    # 2.5 显示待预测样本
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.imshow(image_rgb)
    plt.show()

# 载入项目文件夹
import sys
sys.path.append(r"D:\\Workspace\\MyProjects\\ButterflyClassification")
#################################################################################


#################################################
# 修改者: Xinyu Ou (http://ouxinyu.cn)
# 功能: 使用部署模型对测试集进行评估
# 基本功能：
# 1. 使用部署模型在测试集上进行批量预测，并输出预测结果
# 2. 使用部署模型在测试集上进行单样本预测
#################################################

import numpy as np
import random
import codecs
import os
import cv2
import json
import matplotlib.pyplot as plt
import paddle
from config import project_final_model_path, dataset_root_path, project_result_path
from utils.reader import test_reader
import paddle.nn.functional as F
from utils.dataset import SimplePreprocessing 

# 1. 使用部署模型在测试集上进行批量预测，并输出预测结果
def predict_batch(model, data_reader):
    prediction = []    
    for batch_id, (image, label) in enumerate(data_reader):
        logits = model(image)
        pred = F.softmax(logits)
        pred_id = np.argmax(pred.numpy(), axis=1)
        prediction = np.append(prediction, pred_id).astype('int64')
        
    prediction_path  = os.path.join(project_result_path, 'model_result.txt')
    with codecs.open(prediction_path, 'w', 'utf-8') as f_pred:
        for i in range(len(prediction)):
            f_pred.write('{}\n'.format(prediction[i]))
    print(prediction)
    print('结果文件保存到 {} 成功.'.format(prediction_path)) 

# 2. 使用部署模型在测试集上进行单样本预测
def predict(model, image):    
    isTenCrop = True
    image = SimplePreprocessing(image, isTenCrop=isTenCrop)
    
    if isTenCrop:
        logits = model(image)
        pred = F.softmax(logits)
        pred = np.mean(pred.numpy(), axis=0) 
    else:
        image = paddle.unsqueeze(image, axis=0)
        logits = model(image)
        pred = F.softmax(logits)                
    pred_id = np.argmax(pred)
    
    return pred_id
    
##############################################################
if __name__ == '__main__':
    # 载入模型
    model = paddle.jit.load(project_final_model_path)

    # 1. 计算测试集的准确度 
    predict_batch(model, test_reader())

        # 2. 输出单个样本测试结果
    # 2.1 获取待预测样本的标签信息
    dataset_info_list = os.path.join(dataset_root_path, 'dataset_info.json')
    with open(dataset_info_list, 'r') as f_info:
        dataset_info = json.load(f_info)
    
    # 2.2 从测试列表中随机选择一副图像
    test_list = os.path.join(dataset_root_path, 'test.txt')
    with open(test_list, 'r') as f_test:
        lines = f_test.readlines()
    line = random.choice(lines)
    img_path, label = line.split()
    img_path = os.path.join(dataset_root_path, 'Data', img_path)        
    # img_path = 'D:\\Workspace\\ExpDatasets\\Butterfly\\Data\\zebra\\zeb033.jpg'    
        
    image = cv2.imread(img_path, 1)
    
    # 2.4 给出待测样本的类别
    pred_id = predict(model, image) 
    
    # 将预测的label和ground_turth label转换为label name
    label_name_pred = dataset_info['label_dict'][str(pred_id)]
    print('待测样本的预测类别为：{}'.format(label_name_pred))
    
    # 2.5 显示待预测样本
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.imshow(image_rgb)
    plt.show()