import os
import torch
import timeit
import numpy as np
import pytorch_export_patch
from config import config
from net.nodule_net import NoduleNet
from dataset.bbox_reader import BboxReader
from net.layer.rpn_nms import make_rpn_windows

if __name__ == "__main__": 
    # load model
    os.environ['CUDA_VISIBLE_DEVICES'] = '1' 
    checkpoint = '/fileser/gupc/NoduleNet/results/proxima_data_without_lungmask/model/009.pth'
    model = NoduleNet(config)
    model.eval()
    device = torch.device("cuda")
    model.to(device)
    checkpoint = torch.load(checkpoint)
    model.load_state_dict(checkpoint['state_dict'])
    print("model successfully loaded")


    # # generate data
    # data_dir = config['preprocessed_data_dir']
    # test_set_name = "hahaha"
    # dataset = BboxReader(data_dir, test_set_name, config, mode='train')
    # input_data = dataset[50][0]
    # input_data = input_data.cuda().unsqueeze(0)
    # rpn_windows = make_rpn_windows(device, config)
    # print("data successfully generated")


    # load data
    input_data = torch.from_numpy(np.load('input_data.npy')).cuda()
    rpn_windows = torch.from_numpy(np.load('rpn_windows.npy')).cuda()
    print("data successfully loaded")


    # pth to onnx transform
    model_path = 'NoduleNet.onnx'
    inputs = ['image','windows']
    outputs = ['logits','deltas']
    # dynamic_axes = {'image': {0: 'batch'}, 'bboxes': {0: 'batch'},'labels': {0: 'batch'},'scores': {0: 'batch'}}
    with torch.no_grad():
        torch.onnx.export(model, (input_data,rpn_windows,), model_path, verbose=True, opset_version=10,operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK, input_names=inputs, output_names=outputs) # ,dynamic_axes=dynamic_axes

    # pytorch_export_patch.postprocess_model(model_path)

    


