import os
import json
import mmcv
import time
from mmcv import Config
from mmdet.apis import inference_detector, init_detector, show_result_pyplot, train_detector
from mmdet.models import build_detector
from mmdet.datasets import build_dataset
from mmcv.runner import load_checkpoint
from tqdm import tqdm
import warnings
import torch
import numpy as np
import cv2
import sys
sys.path.append("..")
from ..version import __version__

warnings.filterwarnings("ignore")


class MMDetection:
    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='FasterRCNN',
                 num_classes=-1,
                 dataset_path=None
                 ):

        # 获取外部运行py的绝对路径
        self.cwd = os.path.dirname(os.getcwd())
        # 获取当前文件的绝对路径
        self.file_dirname = os.path.dirname(os.path.abspath(__file__))
        self.save_fold = None
        self.is_sample = False
        self.config = os.path.join(
            self.file_dirname, 'models', 'FasterRCNN/FasterRCNN.py')
        self.checkpoint = os.path.join(
            self.file_dirname, 'models', '/FasterRCNN/FasterRCNN.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 = {
            "FasterRCNN": os.path.join(self.file_dirname, 'models', 'FasterRCNN/FasterRCNN.py'),
            "Yolov3": os.path.join(self.file_dirname, 'models', 'Yolov3/Yolov3.py'),
            "SSD_Lite": os.path.join(self.file_dirname, 'models', 'SSD_Lite/SSD_Lite.py'),
            # "Mask_RCNN":os.path.join(self.file_dirname, 'models', 'Mask_RCNN/Mask_RCNN.py'),
            # 下略
        }
        self.num_classes = num_classes
        self.chinese_res = None
        self.is_sample = False

    def train(self, random_seed=0, save_fold=None, distributed=False, validate=True, device='cpu',
              metric='bbox', save_best='bbox_mAP', optimizer="SGD", epochs=100, lr=0.001, weight_decay=0.001,
              Frozen_stages=1,
              checkpoint=None, batch_size=None):

        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.")
        # 加载网络模型的配置文件
        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/det_model')
            # 如果外部传入save_fold，我们使用传入值
            else:
                self.save_fold = save_fold

        self.cfg.model.backbone.frozen_stages = Frozen_stages

        if self.num_classes != -1:
            if "RCNN" not in self.backbone:  # 单阶段
                self.cfg.model.bbox_head.num_classes = self.num_classes
            elif self.backbone == "FasterRCNN":  # rcnn系列 双阶段
                self.cfg.model.roi_head.bbox_head.num_classes = self.num_classes
            elif self.backbone == "Mask_RCNN":
                self.cfg.model.roi_head.bbox_head.num_classes = self.num_classes
                self.cfg.model.roi_head.mask_head.num_classes = self.num_classes

        self.load_dataset(self.dataset_path)
        # 添加需要进行检测的类名
        if self.backbone in ["Yolov3"]:
            self.cfg.classes = self.get_classes(self.cfg.data.train.dataset.ann_file)
        else:
            self.cfg.classes = self.get_classes(self.cfg.data.train.ann_file)

        # 分别为训练、测试、验证添加类名
        if self.backbone in ["Yolov3"]:
            self.cfg.data.train.dataset.classes = self.cfg.classes
        else:
            self.cfg.data.train.classes = self.cfg.classes
        self.cfg.data.test.classes = self.cfg.classes
        self.cfg.data.val.classes = self.cfg.classes

        # 进行
        self.cfg.work_dir = self.save_fold
        # 创建工作目录
        mmcv.mkdir_or_exist(os.path.abspath(self.cfg.work_dir))
        # 创建分类器
        datasets = [build_dataset(self.cfg.data.train)]
        model = build_detector(self.cfg.model, train_cfg=self.cfg.get(
            'train_cfg'), test_cfg=self.cfg.get('test_cfg'))
        # print("checkpoint", checkpoint)
        if not checkpoint:
            model.init_weights()
        else:
            checkpoint = os.path.abspath(checkpoint)  # pip修改2
            load_checkpoint(model, checkpoint, map_location=torch.device(device))

        model.CLASSES = self.cfg.classes
        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  # 最大的训练轮次

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

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

        meta_info = {
            'tool':'MMEdu'+__version__,
            'task':'Detection',
            '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_detector(
            model,
            datasets,
            self.cfg,
            distributed=distributed,
            validate=validate,
            timestamp=t,
            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"[{'类别标签': 0, '置信度': 1.0, '坐标': {'x': 26, 'y': 81, 'w': 497, 'h': 414}},{'类别标签': 1, '置信度': 1.0, '坐标': {'x1': 250, 'y1': 103, 'x2': 494, 'y2': 341}}]"
            print(sample_result)
        else:
            print("检测结果如下：")
            print(self.chinese_res)
        return self.chinese_res

    def load_checkpoint(self, checkpoint=None, device='cpu',
                        rpn_threshold=0.7, rcnn_threshold=0.7):
        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.")
        print("========= begin inference ==========")
        if self.num_classes != -1 and self.backbone not in ["Yolov3", "SSD_Lite"]:
            self.cfg.model.roi_head.bbox_head.num_classes = self.num_classes

        if checkpoint:
            # 加载数据集及配置文件的路径
            # self.load_dataset(self.dataset_path)
            # 修正检测的目标
            # self.cfg.classes = self.get_class(class_path)
            self.cfg.classes = torch.load(checkpoint, map_location=torch.device('cpu'))['meta']['CLASSES']
            self.cfg.data.train.classes = self.cfg.classes
            self.cfg.data.test.classes = self.cfg.classes
            self.cfg.data.val.classes = self.cfg.classes
            if "RCNN" not in self.backbone:  # 单阶段
                self.cfg.model.bbox_head.num_classes = len(self.cfg.classes)
            else:  # rcnn系列 双阶段
                self.cfg.model.roi_head.bbox_head.num_classes = len(self.cfg.classes)
            # self.cfg.model.roi_head.bbox_head.num_classes = len(self.cfg.classes)
            self.infer_model = init_detector(self.cfg, checkpoint, device=device)
            self.infer_model.CLASSES = self.cfg.classes
        else:
            self.infer_model = init_detector(self.cfg, self.checkpoint, device=device)
        if self.backbone not in ["Yolov3",
                                 "SSD_Lite"]: self.infer_model.test_cfg.rpn.nms.iou_threshold = 1 - rpn_threshold
        if self.backbone not in ["Yolov3",
                                 "SSD_Lite"]: self.infer_model.test_cfg.rcnn.nms.iou_threshold = 1 - rcnn_threshold

    def fast_inference(self, image, show=False, save_fold='det_result'):
        img_array = mmcv.imread(image)
        try:
            self.infer_model
        except:
            print("请先使用load_checkpoint()方法加载权重！")
            return
        result = inference_detector(self.infer_model, img_array)  # 此处的model和外面的无关,纯局部变量
        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 = []
        for i in range(len(result)):
            for j in range(result[i].shape[0]):
                tmp = {}
                tmp['类别标签'] = i
                tmp['置信度'] = result[i][j][4]
                tmp['坐标'] = {"x": int(result[i][j][0]), "y": int(
                    result[i][j][1]), 'w': int(result[i][j][2]), 'h': int(result[i][j][3])}
                # img.append(tmp)
                chinese_res.append(tmp)
        # print(chinese_res)
        self.chinese_res = chinese_res
        # print("========= finish inference ==========")
        return result

    def inference(self, device='cpu',
                  checkpoint=None,
                  image=None,
                  show=True,
                  rpn_threshold=0.7,
                  rcnn_threshold=0.7,
                  save_fold='det_result',
                  ):
        # self.cfg.classes = self.get_class(class_path)
        # self.num_classes = len(self.get_class(class_path))
        # if self.num_classes != -1:
        #     if "RCNN" not in self.backbone: # 单阶段
        #         self.cfg.model.bbox_head.num_classes =self.num_classes
        #     elif self.backbone == "FasterRCNN": # rcnn系列 双阶段
        #         self.cfg.model.roi_head.bbox_head.num_classes = self.num_classes
        #     elif self.backbone == "Mask_RCNN":
        #         self.cfg.model.roi_head.bbox_head.num_classes = self.num_classes
        #         self.cfg.model.roi_head.mask_head.num_classes = self.num_classes

        if type(image)!=np.ndarray and image == None:
            self.is_sample = True
            sample_return = """
        [array([[ 26.547777  ,  81.55447   , 497.37015   , 414.4934    ,
          1.0]], dtype=float32), 
          array([[2.5098564e+02, 1.0334784e+02, 4.9422855e+02, 3.4187744e+02,
        1.0],  dtype=float32)]
            """
            return sample_return
        self.is_sample = False
        print("========= begin inference ==========")

        if self.num_classes != -1 and self.backbone not in ["Yolov3", "SSD_Lite"]:
            self.cfg.model.roi_head.bbox_head.num_classes = self.num_classes

        if checkpoint:
            # 加载数据集及配置文件的路径
            # self.load_dataset(self.dataset_path)
            # 修正检测的目标
            self.cfg.classes = torch.load(checkpoint, map_location=torch.device('cpu'))['meta']['CLASSES']

            self.num_classes = len(self.cfg.classes)
            if self.num_classes != -1:
                if "RCNN" not in self.backbone:  # 单阶段
                    self.cfg.model.bbox_head.num_classes = self.num_classes
                elif self.backbone == "FasterRCNN":  # rcnn系列 双阶段
                    self.cfg.model.roi_head.bbox_head.num_classes = self.num_classes
                elif self.backbone == "Mask_RCNN":
                    self.cfg.model.roi_head.bbox_head.num_classes = self.num_classes
                    self.cfg.model.roi_head.mask_head.num_classes = self.num_classes
            self.cfg.data.train.classes = self.cfg.classes
            self.cfg.data.test.classes = self.cfg.classes
            self.cfg.data.val.classes = self.cfg.classes
            if self.backbone not in ["Yolov3", "SSD_Lite"]:
                self.cfg.model.roi_head.bbox_head.num_classes = len(self.cfg.classes)
            model = init_detector(self.cfg, checkpoint, device=device)
            model.CLASSES = self.cfg.classes
        else:
            model = init_detector(self.cfg, self.checkpoint, device=device)
        # model = build_detector(self.cfg.model, train_cfg=self.cfg.get(
        #     'train_cfg'), test_cfg=self.cfg.get('test_cfg'))
        # if not checkpoint:
        #     model.init_weights()
        # else:
        #     checkpoint = os.path.abspath(checkpoint)  # pip修改2
        #     load_checkpoint(model, checkpoint, map_location=torch.device(device))

        if self.backbone not in ["Yolov3", "SSD_Lite"]: model.test_cfg.rpn.nms.iou_threshold = 1 - rpn_threshold
        if self.backbone not in ["Yolov3", "SSD_Lite"]: model.test_cfg.rcnn.nms.iou_threshold = 1 - rcnn_threshold

        results = []
        # if type(image) == np.ndarray:
            # result = inference_detector(model, image)  # 此处的model和外面的无关,纯局部变量
            # print("res",result)
            # if show == True:
            #     show_result_pyplot(model, image, result)
            #     # 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(result)):
            #         for j in range(result[i].shape[0]):
            #             tmp = {}
            #             tmp['类别标签'] = i
            #             tmp['置信度'] = result[i][j][4]
            #             tmp['坐标'] = {"x1": int(result[i][j][0]), "y1": int(
            #                 result[i][j][1]), 'x2': int(result[i][j][2]), 'y2': int(result[i][j][3])}
            #             chinese_res.append(tmp)
            #     self.chinese_res = chinese_res
        if type(image) == np.ndarray or os.path.isfile(image): # 以路径或numpy数组读入图片
            img_array = mmcv.imread(image)
            print(img_array.shape)
            result = inference_detector(
                model, img_array)  # 此处的model和外面的无关,纯局部变量
            
            chinese_res = []
            for i in range(len(result)):
                for j in range(result[i].shape[0]):
                    tmp = {}
                    tmp['类别标签'] = i
                    tmp['置信度'] = result[i][j][4]
                    tmp['坐标'] = {"x1": int(result[i][j][0]), "y1": int(
                        result[i][j][1]), 'x2': int(result[i][j][2]), 'y2': int(result[i][j][3])}
                    # img.append(tmp)
                    chinese_res.append(tmp)
            # print(chinese_res)
            self.chinese_res = chinese_res
            if show == True:
                show_result_pyplot(model, img_array, result)
                if os.path.isfile(image):
                    model.show_result(img_array, result, show=show, out_file=os.path.join(save_fold, os.path.split(image)[1]))
            print("========= finish inference ==========")
            return result
        else:
            img_dir = image
            mmcv.mkdir_or_exist(os.path.abspath(save_fold))
            chinese_results = []
            for i, img in enumerate(tqdm(os.listdir(img_dir))):
                result = inference_detector(
                    model, img_dir + img)  # 此处的model和外面的无关,纯局部变量
                if show == True:
                    model.show_result(img_dir + img, result,
                                  out_file=os.path.join(save_fold, img))
                chinese_res = []
                for i in range(len(result)):
                    for j in range(result[i].shape[0]):
                        tmp = {}
                        tmp['类别标签'] = i
                        tmp['置信度'] = result[i][j][4]
                        tmp['坐标'] = {"x1": int(result[i][j][0]), "y1": int(
                            result[i][j][1]), 'x2': int(result[i][j][2]), 'y2': int(result[i][j][3])}
                        # img.append(tmp)
                        chinese_res.append(tmp)
                chinese_results.append(chinese_res)
                results.append(result)
            self.chinese_res = chinese_results
        print("========= finish inference ==========")
        return results

    def load_dataset(self, path):
        self.dataset_path = path

        # 数据集修正为coco格式
        if self.backbone in ["Yolov3"]:
            self.cfg.data.train.dataset.img_prefix = os.path.join(self.dataset_path, 'images/train/')
            self.cfg.data.train.dataset.ann_file = os.path.join(self.dataset_path, 'annotations/train.json')
        else:
            self.cfg.data.train.img_prefix = os.path.join(self.dataset_path, 'images/train/')
            self.cfg.data.train.ann_file = os.path.join(self.dataset_path, 'annotations/train.json')

        if os.path.exists(os.path.join(self.dataset_path, 'images/valid/')):
            self.cfg.data.val.img_prefix = os.path.join(self.dataset_path, 'images/valid/')
        else:
            self.cfg.data.val.img_prefix = os.path.join(self.dataset_path, 'images/test/')
        self.cfg.data.val.ann_file = os.path.join(self.dataset_path, 'annotations/valid.json')

        self.cfg.data.test.img_prefix = os.path.join(self.dataset_path, 'images/valid/')
        self.cfg.data.test.ann_file = os.path.join(self.dataset_path, 'annotations/valid.json')

    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 get_classes(self, annotation_file):
        classes = ()
        with open(annotation_file, 'r') as f:
            dataset = json.load(f)
            # categories = dataset["categories"]
            if 'categories' in dataset:
                for cat in dataset['categories']:
                    classes = classes + (cat['name'],)
        return classes

    def convert(self, checkpoint=None, backend="ONNX", out_file="convert_model.onnx", device='cpu'):
        import onnx
        from onnx import load_model
        import os.path as osp
        # from mmdet.core.export import build_model_from_cfg

        ashape = self.cfg.test_pipeline[1].img_scale
        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
        if self.backbone not in ["Yolov3", "SSD_Lite"]:
            self.cfg.model.roi_head.bbox_head.num_classes = self.num_classes
        else:
            self.cfg.model.bbox_head.num_classes = self.num_classes

        # build the model and load checkpoint
        # detector = build_detector(self.cfg.model)
        # model = build_model_from_cfg(self.config, checkpoint)

        if checkpoint:
            # 加载数据集及配置文件的路径
            # self.load_dataset(self.dataset_path)
            # 修正检测的目标
            self.cfg.classes = torch.load(checkpoint, map_location=torch.device('cpu'))['meta']['CLASSES']
            self.num_classes = len(self.cfg.classes)
            if self.num_classes != -1:
                if "RCNN" not in self.backbone:  # 单阶段
                    self.cfg.model.bbox_head.num_classes = self.num_classes
                elif self.backbone == "FasterRCNN":  # rcnn系列 双阶段
                    self.cfg.model.roi_head.bbox_head.num_classes = self.num_classes
                elif self.backbone == "Mask_RCNN":
                    self.cfg.model.roi_head.bbox_head.num_classes = self.num_classes
                    self.cfg.model.roi_head.mask_head.num_classes = self.num_classes
            self.cfg.data.train.classes = self.cfg.classes
            self.cfg.data.test.classes = self.cfg.classes
            self.cfg.data.val.classes = self.cfg.classes
            if self.backbone not in ["Yolov3", "SSD_Lite"]:
                self.cfg.model.roi_head.bbox_head.num_classes = len(self.cfg.classes)
            model = init_detector(self.cfg, checkpoint, device=device)
            model.CLASSES = self.cfg.classes
        else:
            model = init_detector(self.cfg, self.checkpoint, device=device)
        if self.backbone not in ["Yolov3", "SSD_Lite"]: model.test_cfg.rpn.nms.iou_threshold = 0.3  # 1 - rpn_threshold
        if self.backbone not in ["Yolov3",
                                 "SSD_Lite"]: model.test_cfg.rcnn.nms.iou_threshold = 0.3  # 1 - rcnn_threshold

        # detector = build_detector(self.cfg.model, test_cfg=self.cfg.get('test_cfg'))
        # detector.CLASSES = self.num_classes
        normalize_cfg = parse_normalize_cfg(self.cfg.test_pipeline)
        input_img = osp.join(osp.dirname(__file__), './demo/demo.jpg')
        if backend == "ONNX" or backend == 'onnx':
            pytorch2onnx(
                # detector,
                model,
                input_img,
                input_shape,
                normalize_cfg,
                show=False,
                output_file=out_file,
                verify=False,
                test_img=None,
                do_simplify=False)
        else:
            print("Sorry, we only suport ONNX up to now.")
        classes_list = torch.load(checkpoint, map_location=torch.device('cpu'))['meta']['CLASSES']
        onnx_model = load_model(out_file)
        model_info = {'codebase': 'MMDet', '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 parse_normalize_cfg(test_pipeline):
    transforms = None
    for pipeline in test_pipeline:
        if 'transforms' in pipeline:
            transforms = pipeline['transforms']
            break
    assert transforms is not None, 'Failed to find `transforms`'
    norm_config_li = [_ for _ in transforms if _['type'] == 'Normalize']
    assert len(norm_config_li) == 1, '`norm_config` should only have one'
    norm_config = norm_config_li[0]
    return norm_config


def pytorch2onnx(model,
                 input_img,
                 input_shape,
                 normalize_cfg,
                 opset_version=11,
                 show=False,
                 output_file='tmp.onnx',
                 verify=False,
                 test_img=None,
                 do_simplify=False,
                 dynamic_export=None,
                 skip_postprocess=False):
    from mmdet.core.export import build_model_from_cfg, preprocess_example_input
    from mmdet.core.export.model_wrappers import ONNXRuntimeDetector
    from functools import partial
    from mmcv import Config, DictAction
    import onnx
    input_config = {
        'input_shape': input_shape,
        'input_path': input_img,
        'normalize_cfg': normalize_cfg
    }
    # prepare input
    one_img, one_meta = preprocess_example_input(input_config)
    img_list, img_meta_list = [one_img], [[one_meta]]

    if skip_postprocess:
        warnings.warn('Not all models support export onnx without post '
                      'process, especially two stage detectors!')
        model.forward = model.forward_dummy
        torch.onnx.export(
            model,
            one_img,
            output_file,
            input_names=['input'],
            export_params=True,
            keep_initializers_as_inputs=True,
            do_constant_folding=True,
            verbose=show,
            opset_version=opset_version)

        print(f'Successfully exported ONNX model without '
              f'post process: {output_file}')
        return

    # replace original forward function
    origin_forward = model.forward
    model.forward = partial(
        model.forward,
        img_metas=img_meta_list,
        return_loss=False,
        rescale=False)

    output_names = ['dets', 'labels']
    if model.with_mask:
        output_names.append('masks')
    input_name = 'input'
    dynamic_axes = None
    if dynamic_export:
        dynamic_axes = {
            input_name: {
                0: 'batch',
                2: 'height',
                3: 'width'
            },
            'dets': {
                0: 'batch',
                1: 'num_dets',
            },
            'labels': {
                0: 'batch',
                1: 'num_dets',
            },
        }
        if model.with_mask:
            dynamic_axes['masks'] = {0: 'batch', 1: 'num_dets'}

    torch.onnx.export(
        model,
        img_list,
        output_file,
        input_names=[input_name],
        output_names=output_names,
        export_params=True,
        keep_initializers_as_inputs=True,
        do_constant_folding=True,
        verbose=show,
        opset_version=opset_version,
        dynamic_axes=dynamic_axes)

    model.forward = origin_forward

    if do_simplify:
        import onnxsim

        from mmdet 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:
            warnings.warn('Failed to simplify ONNX model.')
    print(f'Successfully exported ONNX model: {output_file}')

    if verify:
        # check by onnx
        onnx_model = onnx.load(output_file)
        onnx.checker.check_model(onnx_model)
        # print(model.CLASSES)
        # wrap onnx model
        onnx_model = ONNXRuntimeDetector(output_file, model.CLASSES, 0)
        if dynamic_export:
            # scale up to valid dynamic shape
            h, w = [int((_ * 1.5) // 32 * 32) for _ in input_shape[2:]]
            h, w = min(1344, h), min(1344, w)
            input_config['input_shape'] = (1, 3, h, w)

        if test_img is None:
            input_config['input_path'] = input_img

        # prepare input once again
        one_img, one_meta = preprocess_example_input(input_config)
        img_list, img_meta_list = [one_img], [[one_meta]]

        # get pytorch output
        with torch.no_grad():
            pytorch_results = model(
                img_list,
                img_metas=img_meta_list,
                return_loss=False,
                rescale=True)[0]

        img_list = [_.cuda().contiguous() for _ in img_list]
        if dynamic_export:
            img_list = img_list + [_.flip(-1).contiguous() for _ in img_list]
            img_meta_list = img_meta_list * 2
        # get onnx output
        onnx_results = onnx_model(
            img_list, img_metas=img_meta_list, return_loss=False)[0]
        # visualize predictions
        score_thr = 0.3
        if show:
            out_file_ort, out_file_pt = None, None
        else:
            out_file_ort, out_file_pt = 'show-ort.png', 'show-pt.png'

        show_img = one_meta['show_img']
        model.show_result(
            show_img,
            pytorch_results,
            score_thr=score_thr,
            show=True,
            win_name='PyTorch',
            out_file=out_file_pt)
        onnx_model.show_result(
            show_img,
            onnx_results,
            score_thr=score_thr,
            show=True,
            win_name='ONNXRuntime',
            out_file=out_file_ort)

        # compare a part of result
        if model.with_mask:
            compare_pairs = list(zip(onnx_results, pytorch_results))
        else:
            compare_pairs = [(onnx_results, pytorch_results)]
        err_msg = 'The numerical values are different between Pytorch' + \
                  ' and ONNX, but it does not necessarily mean the' + \
                  ' exported ONNX model is problematic.'
        # check the numerical value
        for onnx_res, pytorch_res in compare_pairs:
            for o_res, p_res in zip(onnx_res, pytorch_res):
                np.testing.assert_allclose(
                    o_res, p_res, rtol=1e-03, atol=1e-05, err_msg=err_msg)
        print('The numerical values are the same between Pytorch and ONNX')