46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
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)
|