import os
import numpy as np
from neural_compressor.experimental import Quantization, common

model_name = 'conv_asymmetric_pads.onnx'
model_path = os.path.join('/path/to/opencv_extra/testdata/dnn/onnx/models', model_name)
qmodel_path = 'quantized_{}_int8_weights.onnx'.format(model_name[:-5])

config_path = '/path/to/qconv_asympad.yaml'

class RandomDataset:
    def __init__(self):
        self.inputs = [np.random.randn(2, 3, 4).astype(np.float32)]

    def __getitem__(self, idx):
        return self.inputs[idx], 1

    def __len__(self):
        return len(self.inputs)

model = onnx.load(model_path)
quantizer = Quantization()
quantizer.calib_dataloader = common.DataLoader(RandomDataset())
quantizer.model = common.Model(model)
q_model = quantizer()
q_model.save(qmodel_path)


# generate input & output data
import onnxruntime as rt
sess = rt.InferenceSession(qmodel_path, None)
input = np.random.uniform(-1, 1, sess.get_inputs()[0].shape).astype("float32")
output = sess.run([sess.get_outputs()[0].name], {sess.get_inputs()[0].name : input})[0]

print("input has sizes {}".format(input.shape))
input_files = os.path.join("data", "input_" + name)
np.save(input_files, input.data)

print("output has sizes {}".format(output.shape))
output_files = os.path.join("data", "output_" + name)
np.save(output_files, np.ascontiguousarray(output.data))