import ktc
import onnx
from PIL import Image
import numpy as np


## Test Various converters.

## Clone https://github.com/kneron/ConvertorExamples first:
## git clone https://github.com/kneron/ConvertorExamples.git
## cd ConvertorExamples && git lfs pull

## Adjust loggin setting to avoid too many messages before start.
# import logging
# logging.basicConfig(level=logging.INFO)

## Keras to ONNX
# result_m = ktc.onnx_optimizer.keras2onnx_flow('/data1/ConvertorExamples/keras_example/onet-0.417197.hdf5')

## Pytorch to ONNX
# import torch
# import torch.onnx
# # Load the pth saved model
# pth_model = torch.load("/data1/ConvertorExamples/pytorch_example/resnet34.pth", map_location='cpu')
# # Export the model
# dummy_input = torch.randn(1, 3, 224, 224)
# torch.onnx.export(pth_model, dummy_input, '/data1/resnet34.onnx', opset_version=11)
# # Load the exported onnx model as an onnx object
# exported_m = onnx.load('/data1/resnet34.onnx')
# # Optimize the exported onnx object
# result_m = ktc.onnx_optimizer.torch_exported_onnx_flow(exported_m)

## Caffe to ONNX
# result_m = ktc.onnx_optimizer.caffe2onnx_flow('/data1/ConvertorExamples/caffe_example/mobilenetv2.prototxt', '/data1/ConvertorExamples/caffe_example/mobilenetv2.caffemodel')

## TF Lite to ONNX
# result_m = ktc.onnx_optimizer.tflite2onnx_flow('/data1/ConvertorExamples/tflite_example/model_unquant.tflite')

## ONNX Optimization
# optimized_m = ktc.onnx_optimizer.onnx2onnx_flow(result_m, eliminate_tail=True)

## Section 3
km = ktc.ModelConfig(1001, "0001", "520", onnx_path="/workspace/examples/LittleNet/mask_inference_opset11_o2o.onnx")
eval_result = km.evaluate()

def preprocess(input_file):
    image = Image.open(input_file)
    image = image.convert("RGB")
    image = Image.fromarray(np.array(image)[...,::-1])
    img_data = np.array(image.resize((512, 512), Image.BILINEAR)) / 255
    return img_data

input_data = [preprocess("/workspace/examples/LittleNet/pytorch_imgs/Abdullah_0001.png")]
inf_results = ktc.kneron_inference(input_data, onnx_file="/workspace/examples/LittleNet/mask_inference_opset11_o2o.onnx", input_names=["input_1"])
print('Section 3 E2E simulator result:')
print(inf_results)

## Section 4
# Preprocess images and create the input mapping
input_images = [
    preprocess("/workspace/examples/LittleNet/pytorch_imgs/Abdullah_0001.png"),
    preprocess("/workspace/examples/LittleNet/pytorch_imgs/Abdullah_0002.png"),
    preprocess("/workspace/examples/LittleNet/pytorch_imgs/Abdullah_0003.png"),
    preprocess("/workspace/examples/LittleNet/pytorch_imgs/Abdullah_0004.png"),
]
input_mapping = {"input_1": input_images}

# Quantization
bie_path = km.analysis(input_mapping, output_bie = None, threads = 4)

# E2E simulator (fixed point)
radix = ktc.get_radix(input_images)
fixed_results = ktc.kneron_inference(input_data, bie_file=bie_path, input_names=["input_1"], radix=radix)
print('Section 4 E2E simulator result:')
print(fixed_results)

## Section 5
# Batch compile
compile_result = ktc.compile([km],output_dir="/workspace/examples/LittleNet/compile_output_from_matlab")

# E2E simulator (hardware)
hw_results = ktc.kneron_inference(input_data, nef_file=compile_result, radix=radix)
print('Section 5 E2E simulator result:')
print(hw_results)

