import onnx
import ktc
import kneronnxopt
from PIL import Image
import numpy as np
import os
from os import walk

def preprocess(input_file):
    image = Image.open(input_file)
    image = image.convert("RGB")
    img_data = np.array(image.resize((640, 640), Image.BILINEAR)) / 255
    img_data = np.transpose(img_data, (2, 0, 1))
    img_data = np.expand_dims(img_data, 0)
    return img_data


def main():
    
    original_m = onnx.load("YOLOX/yolox_barcode.onnx")
    optimized_m = kneronnxopt.optimize(original_m)
    onnx.save(optimized_m, 'yolox_barcode_save.onnx')
    
    km = ktc.ModelConfig(211, "0001", "730", onnx_model=optimized_m)
    eval_result = km.evaluate()
    
    input_data = [preprocess("YOLOX/datasets/Custom_training/test2017/1.jpg")]
    
    floating_point_inf_results = ktc.kneron_inference(input_data, onnx_file='yolox_barcode_save.onnx', input_names=["images"])
    print(f"floating_point_inf_results = {floating_point_inf_results}")    
    
    raw_images = os.listdir("YOLOX/datasets/Custom_training/test2017/")
    input_images = [preprocess("YOLOX/datasets/Custom_training/test2017/" + image_name) for image_name in raw_images]

    bie_path = km.analysis({"images": input_images}, threads = 8)#, optimize=2)
    print("\nFixed-point analysis done. Saved bie model to '" + str(bie_path) + "'")
    
    #1.5.2. 定点模型推断
    input_data = [preprocess("YOLOX/datasets/Custom_training/test2017/1.jpg")]
    
    fixed_point_inf_results = ktc.kneron_inference(input_data, bie_file=bie_path, input_names=["images"], platform=730)
    print(f"fixed_point_inf_results = {fixed_point_inf_results}")
    
    #1.6.1. 批量编译
    nef_path = ktc.compile([km])
    
    
    #1.6.2. 硬件模拟
    binary_inf_results = ktc.kneron_inference(input_data, nef_file=nef_path, input_names=["images"], platform=730)
    print(f"binary_inf_results = {binary_inf_results}")
    

if __name__ == '__main__':
    main()