import os,sys
import mmcv
import time
import torch
import numpy as np
import cv2
import PIL
import json
from mmcv import Config
from mmcls.apis import inference_model, init_model, show_result_pyplot, train_model, set_random_seed, single_gpu_test
from mmcls.models import build_classifier
from mmcls.datasets import  build_dataloader,build_dataset
from mmcv.runner import load_checkpoint
from mmcv.parallel import collate, scatter
from mmcls.datasets.pipelines import Compose
sys.path.append("..")
from ..version import __version__


def pth_info(checkpoint):
    import torch
    ckpt = torch.load(checkpoint,map_location='cpu')
    meta = ckpt['meta']
    m = list(meta.keys())
    if 'device' in m:
        m.insert(0,m.pop(m.index('device')))
    if 'backbone' in m:
        m.insert(0,m.pop(m.index('backbone')))
    if 'tool' in m:
        m.insert(0,m.pop(m.index('tool')))


    print(checkpoint,"相关信息如下:")
    print("="*(len(checkpoint)+9))
    for i in m:
        print(i,":",meta[i])
    print("="*(len(checkpoint)+9))


class MMClassification:
    def sota():
        pypath = os.path.abspath(__file__)
        father = os.path.dirname(pypath)
        models = os.path.join(father, 'models')
        sota_model = []
        for i in os.listdir(models):
            if i[0] != '_':
                sota_model.append(i)
        return sota_model

    def __init__(
            self,
            backbone='LeNet',
            num_classes=-1,
            dataset_path='../dataset/cls/hand_gray',
            # dataset_type = 'ImageNet'
            **kwargs,
    ):
        if len(kwargs) != 0:
            info = "Error Code: -501. No such parameter: " + next(iter(kwargs.keys()))
            raise Exception(info)
        # 获取外部运行py的绝对路径
        self.cwd = os.path.dirname(os.getcwd())
        # 获取当前文件的绝对路径
        self.file_dirname = os.path.dirname(os.path.abspath(__file__))
        self.save_fold = None
        if backbone not in MMClassification.sota():
            if os.path.exists(backbone): # 传入配置文件
                self.config = backbone
                self.cfg = Config.fromfile(self.config)
                self.backbone = backbone
            else:
                info = "Error Code: -302. No such argument: "+backbone +". Currently "+str(MMClassification.sota())+" is available."
                # print(info)
                raise Exception(info)
        elif backbone in MMClassification.sota():

            self.config = os.path.join(self.file_dirname, 'models', 'LeNet/LeNet.py')
            self.checkpoint = os.path.join(self.file_dirname, 'models', 'LeNet/LeNet.pth')

            self.backbone = backbone
            backbone_path = os.path.join(self.file_dirname, 'models', self.backbone)
            ckpt_cfg_list = list(os.listdir(backbone_path))
            for item in ckpt_cfg_list:
                if item[-1] == 'y' and item[0] != '_':  #pip修改1
                    self.config = os.path.join(backbone_path, item)
                elif item[-1] == 'h':
                    self.checkpoint = os.path.join(backbone_path, item)
                else:
                    #     print("Warning!!! There is an unrecognized file in the backbone folder.")
                    pass

            self.cfg = Config.fromfile(self.config)
        self.dataset_path = dataset_path
        self.lr = None
        self.backbonedict = {
            'MobileNet': os.path.join(self.file_dirname, 'models', 'MobileNet/MobileNet.py'),
            'ResNet18': os.path.join(self.file_dirname, 'models', 'ResNet18/ResNet18.py'),
            'ResNet50': os.path.join(self.file_dirname, 'models', 'ResNet50/ResNet50.py'),
            'LeNet': os.path.join(self.file_dirname, 'models', 'LeNet/LeNet.py'),
            # 下略
        }

        self.num_classes = num_classes
        self.chinese_res = None
        self.is_sample = False
        self.image_type = ""

    def train(self, random_seed=0, save_fold=None, distributed=False, validate=True, device="cpu",
              metric='accuracy', save_best='auto', optimizer="SGD", epochs=100, lr=0.01, weight_decay=0.001,
              checkpoint=None,batch_size=None,**kwargs):
        if len(kwargs) != 0:
            info = "Error Code: -501. No such parameter: " + next(iter(kwargs.keys()))
            raise Exception(info)
        if device not in ['cpu','cuda']:
            info = "Error Code: -301. No such argument: "+ device
            raise Exception(info)
        is_cuda  = torch.cuda.is_available()
        if device == 'cpu' and is_cuda:
            print("You can use  'device=cuda' to accelerate !")
        elif device == 'cuda' and not is_cuda:
            raise Exception("Error Code: -301. Your device doesn't support cuda.")
        if validate not in [True, False]:
            info = "Error Code: -303. No such argument: "+ validate
            raise Exception(info)
        if checkpoint != None and checkpoint.split(".")[-1] != 'pth':
            info = "Error Code: -202. Checkpoint file type error:"+ checkpoint
            raise Exception(info)

        set_random_seed(seed=random_seed)
        # 获取config信息
        if self.backbone.split('.')[-1] == 'py':
            self.cfg = Config.fromfile(self.backbone)
        else:
            self.cfg = Config.fromfile(self.backbonedict[self.backbone])

        # 如果外部不指定save_fold
        if not self.save_fold:
            # 如果外部也没有传入save_fold，我们使用默认路径
            if not save_fold:
                self.save_fold = os.path.join(self.cwd, 'checkpoints/cls_model')
            # 如果外部传入save_fold，我们使用传入值
            else:
                self.save_fold = save_fold

        if self.num_classes != -1:
            if 'num_classes' in self.cfg.model.backbone.keys():
                self.cfg.model.backbone.num_classes = self.num_classes
            else:
                self.cfg.model.head.num_classes = self.num_classes

        self.load_dataset(self.dataset_path)
        if  'val_set' not in os.listdir(self.dataset_path):
            print("Unable to validate during training due to lack of validation set !")
        datasets = None
        try:
            datasets = [build_dataset(self.cfg.data.train)]
        except FileNotFoundError as err:
            if not os.path.exists(self.dataset_path):
                info = "Error Code: -101. No such dataset directory:" + self.dataset_path
            else:
                err = str(err).split(":")[-1]
                info = "Error Code: -201. Dataset file type error.  No such file:"+ err
            raise Exception(info)
        # 进行
        self.cfg.work_dir = self.save_fold
        # 创建工作目录
        mmcv.mkdir_or_exist(os.path.abspath(self.cfg.work_dir))
        # 创建分类器
        model = build_classifier(self.cfg.model)
        if not checkpoint:
            model.init_weights()
        else:
            try:
                load_checkpoint(model, checkpoint, map_location=torch.device(device))
                # model = init_model(self.cfg, checkpoint)
            except FileNotFoundError:
                    info = "Error Code: -102. No such checkpoint file:" + checkpoint
                    raise Exception(info)


        # 添加类别属性以方便可视化
        model.CLASSES = datasets[0].CLASSES

        n_class = len(model.CLASSES)
        if n_class <= 5:
            self.cfg.evaluation.metric_options = {'topk': (1,)}
        else:
            self.cfg.evaluation.metric_options = {'topk': (5,)}
        if optimizer == 'Adam':
            self.cfg.optimizer = dict(type='Adam', lr=lr,betas=(0.9, 0.999),eps=1e-8, weight_decay=0.0001)
        elif optimizer == 'Adagrad':
            self.cfg.optimizer = dict(type='Adagrad',lr=lr, lr_decay=0)
        # 根据输入参数更新config文件
        self.cfg.optimizer.lr = lr  # 学习率
        self.cfg.optimizer.type = optimizer  # 优化器
        self.cfg.optimizer.weight_decay = weight_decay  # 优化器的衰减权重
        self.cfg.evaluation.metric = metric  # 验证指标
        self.cfg.evaluation.save_best = save_best  #
        self.cfg.runner.max_epochs = epochs  # 最大的训练轮次

        # 设置每 10 个训练批次输出一次日志
        self.cfg.log_config.interval = 10
        self.cfg.gpu_ids = range(1)

        self.cfg.seed = random_seed
        self.cfg.device = device
        if batch_size is not None:
            self.cfg.data.samples_per_gpu = batch_size

        meta_info = {
            'tool':'MMEdu'+__version__, 
            'task':'Classification',
            'backbone':self.backbone, 
            'device':device,
            'dataset_size':len(datasets[0]),
            'learning_rate':lr
        }
        t = time.strftime('%Y%m%d_%H%M%S', time.localtime())
        train_model(
            model,
            datasets,
            self.cfg,
            distributed=distributed,
            validate=validate,
            timestamp=t,
            device=device,
            meta=meta_info
        )
        file = open(os.path.join(self.save_fold ,str(t)+".log.json"), 'r') 
        log = []
        import ast 
        for i in file.readlines():
            d = ast.literal_eval(i.rstrip('\n'))
            log.append(d)
        return log

    def print_result(self, res=None):
        if self.is_sample == True:
            print("示例分类结果如下：")
            sample_result = r"[{'标签': 2, '置信度': 1.0, '预测结果': 'scissors'}]"
            print(sample_result)
        else:
            print("分类结果如下：")
            print(self.chinese_res)
        return self.chinese_res

    def load_checkpoint(self,
                  checkpoint=None,
                  device='cpu',
                  **kwargs,
                  ):
        if len(kwargs) != 0:
            info = "Error Code: -501. No such parameter: "+ next(iter(kwargs.keys()))
            raise Exception(info)
        if device not in ['cpu','cuda']:
            info = "Error Code: -301. No such argument: "+ device
            raise Exception(info)
        is_cuda  = torch.cuda.is_available()
        if device == 'cpu' and is_cuda:
            print("You can use  'device=cuda' to accelerate !")
        elif device == 'cuda' and not is_cuda:
            raise Exception("Error Code: -301. Your device doesn't support cuda.")
        if checkpoint != None and checkpoint.split(".")[-1] != 'pth':
            info = "Error Code: -202. Checkpoint file type error:"+ checkpoint
            raise Exception(info)
        # if not checkpoint:
            # checkpoint = os.path.join(self.cwd, 'checkpoints/cls_model/hand_gray/latest.pth')
        self.device = device
        classed_name = torch.load(checkpoint, map_location=torch.device('cpu'))['meta']['CLASSES']
        # classed_name = self.get_class(class_path)
        # self.class_path = class_path
        self.num_classes = len(classed_name)
        if self.num_classes != -1:
            if 'num_classes' in self.cfg.model.backbone.keys():
                self.cfg.model.backbone.num_classes = self.num_classes
            else:
                self.cfg.model.head.num_classes = self.num_classes

        checkpoint = os.path.abspath(checkpoint) # pip修改2
        self.checkpoint = checkpoint
        try:
            self.infer_model = init_model(self.cfg, checkpoint, device=device)
        except FileNotFoundError:
            info = "Error Code: -102. No such checkpoint file:"+ checkpoint
            raise Exception(info)
        self.infer_model.CLASSES = classed_name

    def inference(self, device='cpu',
        checkpoint=None,
        image=None,
        show=True,
        save_fold='cls_result',
        **kwargs,
        ):
        if len(kwargs) != 0:
            info = "Error Code: -501. No such parameter: " + next(iter(kwargs.keys()))
            raise Exception(info)
        if device not in ['cpu','cuda']:
            info = "Error Code: -301. No such argument: "+ device
            raise Exception(info)
        is_cuda  = torch.cuda.is_available()
        if device == 'cpu' and is_cuda:
            print("You can use  'device=cuda' to accelerate !")
        elif device == 'cuda' and not is_cuda:
            raise Exception("Error Code: -301. Your device doesn't support cuda.")
        if type(image)!=np.ndarray and image == None: # 传入图片为空，示例输出
            self.is_sample = True
            sample_return = """
{'pred_label': 2, 'pred_score': 0.9930743, 'pred_class': 'scissors'}
            """
            return sample_return
        self.is_sample = False
        # if not isinstance(image,(str, np.array)):
        # if not isinstance(image,str): # 传入图片格式，仅支持str图片路径
        #     info = "Error Code: -304. No such argument:"+ image+"which is" +type(image)
        #     raise Exception(info)
        if type(image) != PIL.PngImagePlugin.PngImageFile and type(image) != np.ndarray and not os.path.exists(image):
            info = "Error Code: -103. No such file:"+ image
            raise Exception(info)
        if type(image) != PIL.PngImagePlugin.PngImageFile and os.path.isfile(image) and image.split(".")[-1].lower() not in ["png","jpg","jpeg","bmp"]:
            info = "Error Code: -203. File type error:"+ image
            raise Exception(info)

        if checkpoint != None and checkpoint.split(".")[-1] != 'pth':
            info = "Error Code: -202. Checkpoint file type error:"+ checkpoint
            raise Exception(info)
        if not checkpoint:
            checkpoint = os.path.join(self.cwd, 'checkpoints/cls_model/hand_gray/latest.pth')

        checkpoint = os.path.abspath(checkpoint) # pip修改2
        self.load_checkpoint(device= device, checkpoint=os.path.abspath(checkpoint))
        return self.fast_inference(image=image, show=show,save_fold=save_fold, **kwargs)

    def fast_inference(self, image, show=False, save_fold='cls_result',**kwargs):
        if len(kwargs) != 0:
            info = "Error Code: -501. No such parameter: " + next(iter(kwargs.keys()))
            raise Exception(info)
        # img_array = mmcv.imread(image, flag='color')
        try:
            self.infer_model
        except:
            print("请先使用load_checkpoint()方法加载权重！")
            return


        print("========= begin inference ==========")
        classed_name = self.infer_model.CLASSES
        self.num_classes = len(classed_name)

        results = []
        dataset_path = os.getcwd()
        if type(image) == PIL.PngImagePlugin.PngImageFile: # 以PIL读入图片
            self.image_type = "pil"
            image = np.array(image)

        if type(image) == np.ndarray:
            if self.cfg.data.test.pipeline[0]['type'] == 'LoadImageFromFile':
                self.cfg.data.test.pipeline.pop(0)

            if self.image_type != "pil": self.image_type = "numpy"
            print("{} image".format(self.image_type))

            if self.backbone != "LeNet":  # 单张图片 其他网络
                image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
                img_array = mmcv.imread(image, flag='color')
                result = inference_model(self.infer_model, img_array)  # 此处的model和外面的无关,纯局部变量
            else:
                image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
                img_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
                data = dict(img=img_gray)
                test_pipeline = Compose(self.cfg.data.test.pipeline)
                data = test_pipeline(data)
                # data = scatter([data], [self.device])[0]
                data_loader = build_dataloader(
                    [data,data],
                    samples_per_gpu=self.cfg.data.samples_per_gpu,
                    workers_per_gpu=self.cfg.data.workers_per_gpu,
                    shuffle=False,
                    round_up=True)
                result = self.batch_infer(self.infer_model, data_loader)
                ff = classed_name
                pred_class = ff[np.argmax(result[0])] if ff[np.argmax(result[0])][-1:] != "\n" else ff[np.argmax(
                    result[0])][:-1]
                result = {
                    'pred_label': np.argmax(result[0]),
                    'pred_score': result[0][np.argmax(result[0])],
                    'pred_class': pred_class,
                }
            if show== True:
                self.infer_model.show_result(image, result, show=show, out_file=os.path.join(save_fold, "{}img.jpg".format(self.image_type)))
            chinese_res = []
            tmp = {}
            if isinstance(result['pred_label'], np.int64):
                result['pred_label'] = int(result['pred_label'])
            if isinstance(result['pred_score'], np.float32):
                result['pred_score'] = float(result['pred_score'])
            tmp['标签'] = result['pred_label']
            tmp['置信度'] = result['pred_score']
            tmp['预测结果'] = result['pred_class']
            chinese_res.append(tmp)
            self.chinese_res = chinese_res

        elif os.path.isfile(image): # 以路径读入图片
            if self.backbone != "LeNet": # 单张图片 其他网络
                img_array = mmcv.imread(image, flag='color')
                result = inference_model(self.infer_model, img_array)  # 此处的model和外面的无关,纯局部变量
            else: # 单张图片 Lenet
                imagename = image.split("/")[-1]
                # build the dataloader
                f = open("test.txt", 'w')
                f.write(imagename)
                f.write(" 1")
                f.write('\n')
                f.write("no.png 0")
                f.close()
                if not os.path.exists("cache"):
                    os.mkdir('cache')
                import shutil
                if not os.path.exists(os.path.join("cache", imagename)):
                    shutil.copyfile(image, os.path.join("cache", imagename))
                shutil.copyfile(image, os.path.join("cache", "no.png"))
                self.cfg.data.test.data_prefix = os.path.join(dataset_path, 'cache')
                self.cfg.data.test.ann_file = os.path.join(dataset_path, 'test.txt')
                # self.cfg.data.test.classes = os.path.abspath(self.class_path)

                dataset = build_dataset(self.cfg.data.test)
                # the extra round_up data will be removed during gpu/cpu collect
                data_loader = build_dataloader(
                    dataset,
                    samples_per_gpu=self.cfg.data.samples_per_gpu,
                    workers_per_gpu=self.cfg.data.workers_per_gpu,
                    shuffle=False,
                    round_up=True)
                result = self.batch_infer(self.infer_model, data_loader)
                os.remove("test.txt")
                shutil.rmtree("cache")
                ff = classed_name
                pred_class = ff[np.argmax(result[0])] if ff[np.argmax(result[0])][-1:] != "\n" else ff[np.argmax(
                    result[0])][:-1]
                result = {
                    'pred_label': np.argmax(result[0]),
                    'pred_score': result[0][np.argmax(result[0])],
                    'pred_class': pred_class,
                }

            if show == True:
                self.infer_model.show_result(image, result, show=show, out_file=os.path.join(save_fold, os.path.split(image)[1]))
            chinese_res = []
            tmp = {}
            if isinstance(result['pred_label'], np.int64):
                result['pred_label'] = int(result['pred_label'])
            if isinstance(result['pred_score'], np.float32):
                result['pred_score'] = float(result['pred_score'])
            tmp['标签'] = result['pred_label']
            tmp['置信度'] = result['pred_score']
            tmp['预测结果'] = result['pred_class']
            # img.append(tmp)
            chinese_res.append(tmp)
            # print(chinese_res)
            self.chinese_res = chinese_res
            print("\n========= finish inference ==========")
            return result
        else:
            if self.backbone != "LeNet": # 文件夹 其他网络
                f = open("test.txt", 'w')
                for image_name in os.listdir(image):
                    f.write(image_name)
                    f.write(" 1")
                    f.write('\n')
                f.close()
                self.cfg.data.test.data_prefix = image
                self.cfg.data.test.ann_file = os.path.join(dataset_path, 'test.txt')
                # self.cfg.data.test.classes = os.path.abspath(self.class_path)

                dataset = build_dataset(self.cfg.data.test)
                os.remove("test.txt")
                # the extra round_up data will be removed during gpu/cpu collect
                data_loader = build_dataloader(
                    dataset,
                    samples_per_gpu=self.cfg.data.samples_per_gpu,
                    workers_per_gpu=self.cfg.data.workers_per_gpu,
                    shuffle=False,
                    round_up=True)

            else: # 文件夹 Lenet
                dirname = [x.strip() for x in image.split('/') if x.strip() != ''][-1]
                import shutil
                if os.path.exists(os.path.join(dataset_path, 'cache')):
                    shutil.rmtree("cache")
                os.mkdir(os.path.join(dataset_path, 'cache'))
                shutil.copytree(image, os.path.join(dataset_path, 'cache', dirname))
                for i in range(len(classed_name) - 1):
                    dummy_folder = os.path.join(dataset_path, 'cache', 'dummy' + str(i))
                    os.mkdir(dummy_folder)
                self.cfg.data.test.data_prefix = os.path.join(dataset_path, 'cache')
                # self.cfg.data.test.classes = os.path.abspath(self.class_path)
                dataset = build_dataset(self.cfg.data.test)
                # the extra round_up data will be removed during gpu/cpu collect
                data_loader = build_dataloader(
                    dataset,
                    samples_per_gpu=self.cfg.data.samples_per_gpu,
                    workers_per_gpu=self.cfg.data.workers_per_gpu,
                    shuffle=False,
                    round_up=True)

            results_tmp = self.batch_infer(self.infer_model, data_loader)

            if os.path.exists(os.path.join(dataset_path, 'cache')):
                shutil.rmtree("cache")

            results = []
            for i in range(len(results_tmp)):
                pred_class = classed_name[np.argmax(results_tmp[i])] if classed_name[np.argmax(results_tmp[i])][-1:] != "\n" else classed_name[ np.argmax(results_tmp[i])][:-1]
                if isinstance(np.argmax(results_tmp[i]), np.int64):
                    pred_label = int(np.argmax(results_tmp[i]))
                if isinstance(results_tmp[i][np.argmax(results_tmp[i])], np.float32):
                    pred_score = float(results_tmp[i][np.argmax(results_tmp[i])])
                    tmp_result = {
                    'pred_label': pred_label,  # np.argmax(result[i]),
                    'pred_score': pred_score,  # result[i][np.argmax(result[i])],
                    'pred_class': pred_class,
                    }
                    results.append(tmp_result)

            if show == True:
                for i, img in enumerate(os.listdir(image)):
                    self.infer_model.show_result(os.path.join(image,img), results[i], out_file=os.path.join(save_fold, os.path.split(img)[1]))

            # model.show_result(image, result, show=show, out_file=os.path.join(save_fold, os.path.split(image)[1]))
            chinese_res = []
            for i in range(len(results)):
                tmp = {
                    '标签': results[i]['pred_label'],
                    '置信度': results[i]['pred_score'],
                    '预测结果': results[i]['pred_class']
                }
                # img.append(tmp)
                chinese_res.append(tmp)
            # print(chinese_res)
            self.chinese_res = chinese_res
            print("\n========= finish inference ==========")

        return results

    def load_dataset(self, path,**kwargs):
        if len(kwargs) != 0:
            info = "Error Code: -501. No such parameter: "+ next(iter(kwargs.keys()))
            raise Exception(info)
        self.dataset_path = path
        if not isinstance(path, str):
            info = "Error Code: -201. Dataset file type error, which should be <class 'str'> instead of "+ type(path)+"."
            raise Exception(info)
        if not os.path.exists(path):   # 数据集路径不存在
            info = "Error Code: -101. No such dateset directory: "+ path
            raise Exception(info)
        val_set = os.path.join(path, 'val_set')
        val_txt = os.path.join(path, 'val.txt')
        if os.path.exists(val_set) and os.path.exists(val_txt):
            val_num = 0
            for i in os.listdir(val_set):
                val_num += len(os.listdir(os.path.join(val_set,i)))
            if val_num != len(open(val_txt).readlines()):
                info = "Error Code: -201. Dataset file type error. The number of val set images does not match that in val.txt"
                raise Exception(info)
        self.cfg.img_norm_cfg = dict(
            mean=[124.508, 116.050, 106.438],
            std=[58.577, 57.310, 57.437],
            to_rgb=True
        )
        tot = 0
        for root, dirs, files in os.walk(path):
            for file in files:
                if "txt" not in file:
                    impath = os.path.join(root, file)
                    # print("impath", impath)
                    if  "jpg"  in impath or "png" in impath:
                        # print(impath)
                        img = cv2.imread(impath)
                        try:
                            img.shape
                        except AttributeError:
                            info = "Error Code: -201. The image file {} is damaged.".format(impath)
                            raise Exception(info)

                        tot += 1
        val_permit = True
        if 'val.txt' not in os.listdir(path) and 'val_set' in os.listdir(path):
            # print("生成val.txt")
            val_permit, valtxt_path = self.generate_txt(path, "val")
        # test_permit = True
        # if 'test.txt' not in os.listdir(path) and 'test_set' in os.listdir(path):
        #     print("生成test.txt")
        #     test_permit, testtxt_path = self.generate_txt(path, "test")
        class_permit = True
        if 'classes.txt' not in os.listdir(path):
            # print("生成classes.txt")
            training_set = os.path.join(path, 'training_set')
            content = sorted(os.listdir(training_set))
            content = [i+"\n" for i in content]
            try:
                classestxt = open(os.path.join(path,"classes.txt"), mode='w')
            except:
                class_permit = False
                dataset_txt = "dataset_txt"
                if not os.path.exists(dataset_txt):
                    os.mkdir(dataset_txt)
                classtxt_path = os.path.join(dataset_txt,"classes.txt")
                classestxt = open(classtxt_path, mode='w')

            classestxt.writelines(content)
            classestxt.close()

        self.cfg.data.train.data_prefix = os.path.join(self.dataset_path, 'training_set')
        # self.cfg.data.train.classes = os.path.join(self.dataset_path, 'classes.txt')

        self.cfg.data.val.data_prefix = os.path.join(self.dataset_path, 'val_set')
        if val_permit:
            self.cfg.data.val.ann_file = os.path.join(self.dataset_path, 'val.txt')
        else:
            self.cfg.data.val.ann_file = valtxt_path
        # self.cfg.data.val.classes = os.path.join(self.dataset_path, 'classes.txt')

        self.cfg.data.test.data_prefix = os.path.join(self.dataset_path, 'test_set')
        # if test_permit:
        self.cfg.data.test.ann_file = os.path.join(self.dataset_path, 'test.txt')
        # else:
        #     self.cfg.data.test.ann_file = testtxt_path
        # self.cfg.data.test.classes = os.path.join(self.dataset_path, 'classes.txt')
        if class_permit:
            self.cfg.data.train.classes = os.path.join(self.dataset_path, 'classes.txt')
            self.cfg.data.val.classes = os.path.join(self.dataset_path, 'classes.txt')
            self.cfg.data.test.classes = os.path.join(self.dataset_path, 'classes.txt')
        else:
            self.cfg.data.train.classes = classtxt_path
            self.cfg.data.val.classes = classtxt_path
            self.cfg.data.test.classes = classtxt_path

    def generate_txt(self, path, type):
        permit = True
        val_set = os.path.join(path, type+'_set')
        txt_path = os.path.join(path,type+".txt")
        try: 
            valtxt = open(txt_path, mode='w')
        except:
            permit = False
            dataset_txt = "dataset_txt"
            if not os.path.exists(dataset_txt):
                os.mkdir(dataset_txt)
            txt_path = os.path.join(dataset_txt,type+".txt")
            valtxt = open(txt_path, mode='w')

        content = []
        for label, i in enumerate(sorted(os.listdir(val_set))):
            for j in sorted(os.listdir(os.path.join(val_set,i))):
                # print(os.path.join(i,j), label)
                content.append("{} {}\n".format(os.path.join(i,j), label))
        valtxt.writelines(content)
        valtxt.close()
        return permit, txt_path

    def get_class(self, class_path):
        classes = []
        with open(class_path, 'r') as f:
            for name in f:
                classes.append(name.strip('\n'))
        return classes

    def batch_infer(self, model, data_loader):
        results_tmp = []
        model.eval()
        results = []
        dataset = data_loader.dataset
        prog_bar = mmcv.ProgressBar(task_num=len(dataset),start=False)
        from mmcv.utils.timer import Timer
        prog_bar.file.flush()
        prog_bar.timer = Timer()
        for i, data in enumerate(data_loader):
            # data = data.to(device)
            if self.device == "cuda": data = scatter(data, [self.device])[0]
            with torch.no_grad():
                result = model(return_loss=False, **data)

            batch_size = len(result)
            results_tmp.extend(result)

            batch_size = data['img'].size(0)
            for _ in range(batch_size):
                prog_bar.update()
        return results_tmp
    
    def convert(self, checkpoint=None, backend="ONNX", out_file="convert_model.onnx"):
        import onnx
        from onnx import load_model
        if not (backend == "ONNX" or backend == 'onnx'):
            print("Sorry, we only suport ONNX up to now.")
            return
        state_dict = torch.load(checkpoint, map_location=torch.device('cpu'))
        classes_list = state_dict['meta']['CLASSES']
        self.num_classes = len(classes_list)
        if self.backbone == 'LeNet':
            from mmcls.models.backbones import LeNet5
            from collections import OrderedDict
            model = LeNet5(num_classes=self.num_classes)

            class LeNet5_SoftMax(LeNet5):
                def forward(self, x):
                    x = self.features(x)
                    if self.num_classes > 0:
                        x = self.classifier(x.squeeze())
                        x = torch.softmax(x, dim=0)
                    return (x,)

            model = LeNet5_SoftMax(num_classes=self.num_classes)
            new_state_dict = OrderedDict()
            for key in state_dict['state_dict']:
                new_state_dict[key[9:]] = state_dict['state_dict'][key]
            model.load_state_dict(new_state_dict)
            dummy_input = torch.randn(1, 1, 32, 32)
            try:
                torch.onnx.export(model, dummy_input, out_file)
                print(f'Successfully exported ONNX model: {out_file}')
            except:
                print('Please use the checkpoint train by MMEdu')
        else:
            ashape = [224, 224]
            if len(ashape) == 1:
                input_shape = (1, 3, ashape[0], ashape[0])
            elif len(ashape) == 2:
                input_shape = (
                                  1,
                                  3,
                              ) + tuple(ashape)
            else:
                raise ValueError('invalid input shape')
            self.cfg.model.pretrained = None

            self.cfg.model.head.num_classes = self.num_classes
            # build the model and load checkpoint
            classifier = build_classifier(self.cfg.model)

            if checkpoint:
                load_checkpoint(classifier, checkpoint, map_location='cpu')
            else:
                load_checkpoint(classifier, self.checkpoint, map_location='cpu')

            pytorch2onnx(
                classifier,  # 模型，此处是分类器
                input_shape,
                output_file=out_file,
                do_simplify=False,
                verify=False)

        onnx_model = load_model(out_file)
        model_info = {'codebase': 'MMCls', 'modelname': self.backbone, 'classes':classes_list}

        unicode_string = json.dumps(model_info)
        model_info_metadata = onnx.StringStringEntryProto(key='MODEL_INFO', value=unicode_string)
        
        onnx_model.metadata_props.append(model_info_metadata)
        
        inputs = onnx_model.graph.input
        name_to_input = {}
        for input in inputs:
            name_to_input[input.name] = input

        for initializer in onnx_model.graph.initializer:
            if initializer.name in name_to_input:
                inputs.remove(name_to_input[initializer.name])

        os.remove(out_file)
        onnx.save(onnx_model, out_file)
        with open(out_file.replace(".onnx", ".py"), "w+") as f:

            gen0 = """
import cv2
import BaseDeploy as bd
model_path = '
"""
            gen1 = """'
cap = cv2.VideoCapture(0)
ret, img = cap.read()
model = bd(model_path)
result = model.inference(img)
print(result)
cap.release()
            """

            gen = gen0.strip('\n') + out_file + gen1.strip('\n')
            f.write(gen)

# 模型部署
def _demo_mm_inputs(input_shape, num_classes):
    """Create a superset of inputs needed to run test or train batches.
    Args:
        input_shape (tuple):
            input batch dimensions
        num_classes (int):
            number of semantic classes
    """
    (N, C, H, W) = input_shape
    rng = np.random.RandomState(0)
    imgs = rng.rand(*input_shape)
    gt_labels = rng.randint(
        low=0, high=num_classes, size=(N, 1)).astype(np.uint8)
    mm_inputs = {
        'imgs': torch.FloatTensor(imgs).requires_grad_(True),
        'gt_labels': torch.LongTensor(gt_labels),
    }
    return mm_inputs

def pytorch2onnx(model,
                input_shape,
                opset_version=11,
                dynamic_export=False,
                show=False,
                output_file='tmp.onnx',
                do_simplify=False,
                verify=False):
    """Export Pytorch model to ONNX model and verify the outputs are same
    between Pytorch and ONNX.
    Args:
        model (nn.Module): Pytorch model we want to export.
        input_shape (tuple): Use this input shape to construct
            the corresponding dummy input and execute the model.
        opset_version (int): The onnx op version. Default: 11.
        show (bool): Whether print the computation graph. Default: False.
        output_file (string): The path to where we store the output ONNX model.
            Default: `tmp.onnx`.
        verify (bool): Whether compare the outputs between Pytorch and ONNX.
            Default: False.
    """
    from functools import partial
    import onnxruntime as rt
    from mmcv.onnx import register_extra_symbolics
    model.cpu().eval()

    if hasattr(model.head, 'num_classes'):
        num_classes = model.head.num_classes
    # Some backbones use `num_classes=-1` to disable top classifier.
    elif getattr(model.backbone, 'num_classes', -1) > 0:
        num_classes = model.backbone.num_classes
    else:
        raise AttributeError('Cannot find "num_classes" in both head and '
                            'backbone, please check the config file.')
    mm_inputs = _demo_mm_inputs(input_shape, num_classes)

    imgs = mm_inputs.pop('imgs')
    img_list = [img[None, :] for img in imgs]

    # replace original forward function
    origin_forward = model.forward
    model.forward = partial(model.forward, img_metas={}, return_loss=False)
    register_extra_symbolics(opset_version)

    # support dynamic shape export
    if dynamic_export:
        dynamic_axes = {
            'input': {
                0: 'batch',
                2: 'width',
                3: 'height'
            },
            'probs': {
                0: 'batch'
            }
        }
    else:
        dynamic_axes = {}

    with torch.no_grad():
        torch.onnx.export(
            model, (img_list, ),
            output_file,
            input_names=['input'],
            output_names=['probs'],
            export_params=True,
            keep_initializers_as_inputs=True,
            dynamic_axes=dynamic_axes,
            verbose=show,
            opset_version=opset_version)
        print(f'Successfully exported ONNX model: {output_file}')
    model.forward = origin_forward

    if do_simplify:
        import onnx
        import onnxsim
        from mmcv import digit_version

        min_required_version = '0.4.0'
        assert digit_version(onnxsim.__version__) >= digit_version(
            min_required_version
        ), f'Requires to install onnxsim>={min_required_version}'

        model_opt, check_ok = onnxsim.simplify(output_file)
        if check_ok:
            onnx.save(model_opt, output_file)
            print(f'Successfully simplified ONNX model: {output_file}')
        else:
            print('Failed to simplify ONNX model.')
    if verify:
        # check by onnx
        import onnx
        onnx_model = onnx.load(output_file)
        onnx.checker.check_model(onnx_model)

        # test the dynamic model
        if dynamic_export:
            dynamic_test_inputs = _demo_mm_inputs(
                (input_shape[0], input_shape[1], input_shape[2] * 2,
                input_shape[3] * 2), model.head.num_classes)
            imgs = dynamic_test_inputs.pop('imgs')
            img_list = [img[None, :] for img in imgs]

        # check the numerical value
        # get pytorch output
        pytorch_result = model(img_list, img_metas={}, return_loss=False)[0]

        # get onnx output
        input_all = [node.name for node in onnx_model.graph.input]
        input_initializer = [
            node.name for node in onnx_model.graph.initializer
        ]
        net_feed_input = list(set(input_all) - set(input_initializer))
        assert (len(net_feed_input) == 1)
        sess = rt.InferenceSession(output_file)
        onnx_result = sess.run(
            None, {net_feed_input[0]: img_list[0].detach().numpy()})[0]
        if not np.allclose(pytorch_result, onnx_result):
            raise ValueError(
                'The outputs are different between Pytorch and ONNX')
        print('The outputs are same between Pytorch and ONNX')