# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the BSD 3-Clause License  (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://opensource.org/licenses/BSD-3-Clause
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import argparse
import torch
import torch.onnx
from collections import OrderedDict
sys.path.append(r"./HRNet-Image-Classification")
sys.path.append(r"./HRNet-Image-Classification/lib")
from lib.models import cls_hrnet
from lib.config import config
from lib.config import update_config
def parse_args():
    parser = argparse.ArgumentParser(description='Train classification network')
    parser.add_argument('--cfg',
                        help='experiment configure file name',
                        required=True,
                        type=str)
                        
    parser.add_argument('--modelDir',
                        help='model directory',
                        type=str,
                        default='')
    parser.add_argument('--logDir',
                        help='log directory',
                        type=str,
                        default='')
    parser.add_argument('--dataDir',
                        help='data directory',
                        type=str,
                        default='')
    parser.add_argument('--testModel',
                        help='testModel',
                        type=str,
                        default='')

    parser.add_argument('--input',
                        help='input pytorch model',
                        required=True,
                        type=str)
    
    parser.add_argument('--output',
                        help='output onnx model',
                        required=True,
                        type=str)

    args = parser.parse_args()
    update_config(config, args)

    return args

def proc_node_module(checkpoint, AttrName):
    new_state_dict = OrderedDict()
    for k, v in checkpoint[AttrName].items():
        if(k[0:7] == "module."):
            name = k[7:]
        else:
            name = k[0:]
        new_state_dict[name] = v
    return new_state_dict

def pth2onnx():
    args = parse_args()
    
    modelpth = args.input
    checkpoint = torch.load(modelpth, map_location='cpu')
    model = cls_hrnet.get_cls_net(config)
    output_file = args.output
    
    model.load_state_dict(checkpoint)
    model.eval()
    
    input_names = ["image"]
    output_names = ["class"]
    dynamic_axes = {'image': {0: '-1'}, 'class': {0: '-1'}}
    dummy_input = torch.randn(1, 3, 224, 224)
    torch.onnx.export(model, dummy_input, output_file, input_names = input_names, dynamic_axes = dynamic_axes, output_names = output_names, verbose=True, opset_version=11)

if __name__ == "__main__":
    pth2onnx()