from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import sys
CENTERNET_PATH = '../src'
sys.path.insert(0, CENTERNET_PATH)
MODEL_PATH = '../models/ctdet_coco_dla_2x.pth'
import os
import _init_paths
import torch
import torch.utils.data
from opts import opts
from models.model import create_model, load_model, save_model
from models.data_parallel import DataParallel
from logger import Logger
from datasets.dataset_factory import get_dataset


def convert():
    #device = torch.device("cpu")
    device = torch.device("cuda")
    torch.set_default_tensor_type(torch.cuda.FloatTensor)
    TASK = 'ctdet' 
    opt = opts().parse('{} --load_model {}'.format(TASK, MODEL_PATH).split(' '))
    Dataset = get_dataset(opt.dataset, opt.task)
    opt = opts().update_dataset_info_and_set_heads(opt, Dataset)
    model = create_model(opt.arch, opt.heads, opt.head_conv) 
    model = load_model(model, input_file, None, opt.resume, opt.lr, opt.lr_step)
    model.eval()
    
    input_names = ["actual_input"]
    output_names = ["output1","output2","output3"]
    dynamic_axes = {'actual_input': {0: '-1'}, 'output1': {0: '-1'}, 'output2': {0: '-1'}, 'output3': {0: '-1'}}
    dummy_input = torch.randn(1, 3, 512, 512)
    torch.onnx.export(model, dummy_input, output_file, input_names = input_names, dynamic_axes = dynamic_axes, output_names = output_names, opset_version=11, verbose=True)

if __name__ == "__main__":
    input_file = sys.argv[1]
    output_file = sys.argv[2]
    convert()