import sys
from models import build_model
sys.path.extend(['..','.' ])
from segutils.trtUtils2 import pth2onnx,onnx2engine,onnx_inference
from engine import DictToObject
from pathlib import Path
import torch
import os
import tensorrt as trt
import numpy as np
import argparse


def main(opt):
    #pth_model='../weights/best_mae.pth'
    pth_model=opt.weights.strip()
    onnx_name = pth_model.replace('.pth','_dynamic.onnx')
    trt_name = onnx_name.replace('.onnx','.engine')
    dynamic_hw ={'input':{0:'batch',2:'H',3:'W'},
        'output0':{1:'C'},
        'output1':{1:'C'},
        
       }

        
        
    par = {'backbone':'vgg16_bn', 'gpu_id':0, 'line':2, 'output_dir':'./output', 'row':2, 'anchorFlag':False,'weight_path':'./weights/best_mae.pth'}
    args = DictToObject(par)
    model = build_model(args) 
    
    
    pthFile = Path(pth_model)
    checkpoint = torch.load(pthFile, map_location='cpu')
    model.load_state_dict(checkpoint['model'])
    model = model.to('cuda:0')

  
    inputShape =(1, 3, 128*4,128*4)#(bs,channels,height,width)

    input_profile_shapes = [(1,3,256,256),(1,3,1024,1024),(1,3,2048,2048)]
    
    pth2onnx(model,onnx_name,input_shape=(1,3,512,512),input_names=['input'],output_names=[ 'output0' ,'output1'],dynamix_axis=dynamic_hw)
    
    onnx2engine(onnx_name,trt_name,input_shape=[1,3,-1,-1],half=True,max_batch_size=1,input_profile_shapes=input_profile_shapes)
    
   
    
    

    
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='/mnt/thsw2/DSP2/weights/cityMangement2/weights/urbanManagement/DMPR/dp_detector_499.pth', help='model path(s)')
    opt = parser.parse_args()
    
    main(opt)