from torch.autograd import Variable
import torch
import sys
import noox
from onnx2keras import onnx_to_keras
import json

def main():
    modelfile=sys.argv[1]
    outputfolder=sys.argv[2]
    print(modelfile,outputfolder)
    trained_model=torch.load(modelfile)
    inputsize=list(trained_model.parameters())[0].shape
    dummy_input=Variable(torch.randn(inputsize))
#    dummy_input = Variable(torch.randn(1, 1, 28, 28)) # one black and white 28 x 28 picture will be the input to the model
    torch.onnx.export(trained_model, dummy_input, outputfolder+"/dummy.onnx")
    onnx_model=onnx.load(outputfolder+"/dummy.onnx")
    k_model=onnx_to_keras(onnx_model, ['input'])
    jsonfile=k_model.to_json()
    out_file=open(outputfolder+"/dummy.json","w")
    json.dump(jsonfile,out_file,indent=6)
    
    
    








if __name__ == '__main__':
    main()
