import os
import numpy as np
import torch
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 = []

# 遍歷 voc_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)

        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("\nFixed-point analysis done. Saved bie model to '" + str(bie_model_path) + "'")