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

onnx_model_path = "/Desktop/Kneron/SW/combo_for_kneron.onnx"
optimized_onnx_model_path = '/Desktop/Kneron/SW/optimized_eliminate_tail.onnx'
test_images_path = "/Desktop/Kneron/SW/test_images"


# #########################  
# Export and Optimize model
# #########################
exported_m = onnx.load(onnx_model_path)
result_m = ktc.onnx_optimizer.torch_exported_onnx_flow(exported_m)
optimized_m = ktc.onnx_optimizer.onnx2onnx_flow(result_m, eliminate_tail=True, opt_matmul=False)
onnx.save(optimized_m, optimized_onnx_model_path)


#########################  
#      Evaluate Model
#########################
model_id = 77777
model_version = '0001'
platform = '720'
km = ktc.ModelConfig(model_id, model_version, platform, onnx_model=None, 
                                        onnx_path=optimized_onnx_model_path, 
                                        bie_path=None)
eval_result = km.evaluate()


def preprocess(input_file):
    image = Image.open(input_file)
    image = image.convert("RGB")
    image.save("/Desktop/Kneron/SW/img.png")
    img_data = np.array(image.resize((320, 320), Image.BILINEAR)) / 255
    # img_data = np.transpose(img_data, (1, 0, 2))
    return img_data


# # #########################  
# #       Quantize
# # #########################
# # Preprocess images and create the input mapping
input_images = [preprocess(im) for im in glob(test_images_path + "/*.PNG")]
input_mapping = {"input.1": input_images[:1]}

bie_path = km.analysis(input_mapping, threads = 4)

#########################  
#   Compile to NEF
#########################

# km = ktc.ModelConfig(model_id, model_version, platform, bie_path=bie_path)
model_list = [km]
compile_result = ktc.compile(model_list)