import json
import torch
from PIL import Image
from ecos_core.indad.indad import Model
from easydict import EasyDict
from ecos_core.indad.sample_code.data_transform import transform_default


def convert_to_onnx(opt_path: str, weight_path: str, output_path: str, dummy_input_path:str):
        """Convert INDAD to ONNX

        Args:
            opt_path (str): Path to opt.json
            weight_path (str): path to weight .pt
            output_path (str): ouput path to save onnx
            dummy_input_path (str): dummy input to convert
        """
        # load model
        with open(opt_path, 'r') as f:
            data = EasyDict(json.load(f))
        model = Model(opt=data)
        model_instance = model.build_model(weight_path)
        model_instance.eval()
        # init dummy
        transform = transform_default
        img = Image.open(dummy_input_path)
        img = img.convert("RGB")
        img_trans = transform(img).unsqueeze(dim=0)

        print(model_instance)

        #quantized_model = quantize_dynamic(model_instance, {torch.nn.Module}, dtype=torch.qint8)
        # export onnx
        torch.onnx.export(model_instance, img_trans, output_path,
                        export_params=True,
                        opset_version=14,
                        do_constant_folding=True, 
                        
                        input_names=['input'],
                        output_names=['img_level_score', "pxl_level_score"])