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

kneron_platform = "520"
model_path = "/workspace/ai_training/detection/yolov5/yolov5/yolov5s-noupsample-coco128.onnx"
images_path = "/workspace/ai_training/detection/yolov5/coco128/images/train2017"

# Load the exported onnx model as an onnx object
exported_m = onnx.load(model_path)
# Optimize the exported onnx object
optmized_m = ktc.onnx_optimizer.torch_exported_onnx_flow(exported_m)
optmized_m = kneronnxopt.optimize(optmized_m)

# npu (only) performance simulation
km = ktc.ModelConfig(20008, "0001", kneron_platform, onnx_model=optmized_m)
eval_result = km.evaluate()
print(f"\nNPU performance evaluation result:\n{str(eval_result)}")

img_list = []
for (dirpath, dirnames, filenames) in walk(images_path):
    for f in filenames:
        fullpath = os.path.join(dirpath, f)
        
        image = Image.open(fullpath)
        image = image.convert("RGB")
        image = Image.fromarray(np.array(image)[...,::-1])
        img_data = np.array(image.resize((640, 640), Image.BILINEAR)) / 255
        img_data = img_data.transpose(2, 0 ,1)[None]
        img_list.append(img_data)

# fixed-point analysis
print("\nGenerating .bie model...")
bie_model_path = km.analysis({"images": img_list}, threads=4)
print(f"\nFixed-point analysis done. Saved bie model to '{str(bie_model_path)}'")

# compile
print("\nGenerating .nef model...")
nef_model_path = ktc.compile([km])
print(f"\nCompile done. Saved nef model to '{str(nef_model_path)}'")