import os
import numpy as np
import torch
import ktc  # Kneron Toolchain
from yolov5_preprocess import Yolov5_preprocess  # 使用你的預處理
import kneron_preprocessing

# 設定裝置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 設定圖片大小（與訓練時一致）
imgsz_h, imgsz_w = 640, 640

# 量化數據集目錄（請確保這個資料夾存在）
data_path = "/data50"
img_list = []

# 設定 ONNX 模型路徑（確保這個路徑在 Docker 內部是否正確）
onnx_model_path = "/workspace/yolov5/latest.opt.onnx"

# **初始化 Kneron ModelConfig 物件**
km = ktc.ModelConfig(20008, "0001", "720", onnx_model=onnx_model_path)

# 遍歷 data50 並進行預處理
for root, _, files in os.walk(data_path):
    for f in files:
        fullpath = os.path.join(root, f)
        
        # 執行與訓練相同的預處理
        img_data, _ = Yolov5_preprocess(fullpath, device, imgsz_h, imgsz_w)

        # 確保格式為 NumPy 陣列
        img_data = img_data.cpu().numpy()

        print(f"Processed: {fullpath}")
        img_list.append(img_data)

# 轉為 NumPy 格式
img_list = np.array(img_list)

# **執行 BIE 量化分析**
bie_model_path = km.analysis({"input": img_list})

# 輸出成功訊息
print("\n✅ Fixed-point analysis done! BIE model saved to:", bie_model_path)