import sys
import torch
sys.path.append('models/')
import argparse
import cv2
import numpy as np
import torch
import os

def parse_args():
    """Parses command line arguments."""
    parser = argparse.ArgumentParser(description="export to onnx")
    parser.add_argument("checkpoint", type=str)
    args = parser.parse_args()
    return args


def export_onnx():
    output_model_dir='app/src/main/assets/'
    INPUT_SIZE=224
    #filename='enet_b0_8_best_vgaf'
    filename='enet_b0_8_va_mtl'
    
    example = torch.rand(1, 3, INPUT_SIZE, INPUT_SIZE)
    model=torch.load('models/emotion/'+filename+'.pt').cpu()
    model.eval()
    input_size = [1, 3, 224, 224]
    input_image = (np.random.random(input_size) * 255.)
    input_image = torch.from_numpy(input_image.astype('float32'))
    #output = model(input_image)

    torch.onnx.export(model,
                      example,
                      "emotion.onnx",
                      verbose=False,
                      opset_version=10,
                      do_constant_folding=False,
                      input_names=["input"])

if __name__ == "__main__":
    export_onnx()