111 lines
3.9 KiB
Python
111 lines
3.9 KiB
Python
import ktc
|
|
import numpy as np
|
|
import os
|
|
import onnx
|
|
import shutil
|
|
from PIL import Image
|
|
import torch
|
|
from yolov5_preprocess import Yolov5_preprocess
|
|
import kneron_preprocessing
|
|
|
|
# 設定 ONNX 模型儲存路徑
|
|
onnx_dir = 'runs/train/exp24/weights/'
|
|
onnx_path = os.path.join(onnx_dir, 'best_no_sigmoid.onnx')
|
|
|
|
# 確保目標資料夾存在
|
|
os.makedirs(onnx_dir, exist_ok=True)
|
|
|
|
# 加載並優化 ONNX 模型
|
|
m = onnx.load(onnx_path)
|
|
m = ktc.onnx_optimizer.onnx2onnx_flow(m)
|
|
opt_onnx_path = os.path.join(onnx_dir, 'latest.opt.onnx')
|
|
onnx.save(m, opt_onnx_path)
|
|
|
|
km = ktc.ModelConfig(20008, "0001", "520", onnx_model=m)
|
|
eval_result = km.evaluate()
|
|
print("\nNpu performance evaluation result:\n" + str(eval_result))
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
imgsz_h, imgsz_w = 640, 640
|
|
|
|
data_path = "datacoin"
|
|
files_found = [f for _, _, files in os.walk(data_path) for f in files if f.lower().endswith((".jpg", ".jpeg", ".png", ".bmp"))]
|
|
|
|
if not files_found:
|
|
raise FileNotFoundError(f"\u274c Error: No images found in {data_path}! Please check your dataset.")
|
|
|
|
print(f"\u2705 Found {len(files_found)} images in {data_path}")
|
|
|
|
# 獲取 ONNX 模型的輸入名稱
|
|
input_name = m.graph.input[0].name # 確保 key 與 ONNX input name 一致
|
|
|
|
# 存儲預處理後的圖片數據
|
|
img_list = []
|
|
|
|
# 遍歷 data50 並進行預處理
|
|
for root, _, files in os.walk(data_path):
|
|
for f in files:
|
|
fullpath = os.path.join(root, f)
|
|
|
|
# 只處理圖片文件
|
|
if not f.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
|
|
print(f"\u26a0\ufe0f Skipping non-image file: {fullpath}")
|
|
continue
|
|
|
|
# 嘗試處理圖片
|
|
try:
|
|
img_data, _ = Yolov5_preprocess(fullpath, device, imgsz_h, imgsz_w)
|
|
img_data = img_data.cpu().numpy()
|
|
print(f"\u2705 Processed: {fullpath}")
|
|
img_list.append(img_data)
|
|
except Exception as e:
|
|
print(f"\u274c Failed to process {fullpath}: {e}")
|
|
|
|
# 確保 img_list 不是空的
|
|
if not img_list:
|
|
raise ValueError("\u274c Error: No valid images were processed! Please check the image paths and formats.")
|
|
|
|
# 執行 BIE 量化
|
|
bie_model_path = km.analysis({input_name: img_list})
|
|
|
|
# 確保 BIE 檔案儲存到指定目錄
|
|
bie_save_path = os.path.join(onnx_dir, os.path.basename(bie_model_path))
|
|
shutil.copy(bie_model_path, bie_save_path) # 使用 shutil.move 來處理跨磁碟移動
|
|
|
|
# 確認 BIE 模型是否生成
|
|
if not os.path.exists(bie_save_path):
|
|
raise RuntimeError(f"\u274c Error: BIE model was not generated! Please check your quantization process.")
|
|
|
|
print("\n\u2705 Fixed-point analysis done! BIE model saved to:", bie_save_path)
|
|
|
|
# 確保 `km` 已經初始化，並且 `.bie` 模型已生成
|
|
nef_model_path = ktc.compile([km])
|
|
|
|
# 確保 nef_model_path 不是 None 或空值
|
|
if not nef_model_path:
|
|
raise RuntimeError("❌ Error: ktc.compile() did not return a valid .nef file path!")
|
|
|
|
# 確保 NEF 目標資料夾存在
|
|
os.makedirs(onnx_dir, exist_ok=True)
|
|
|
|
# 確保 nef_model_path 不是 None 或空值
|
|
if not nef_model_path:
|
|
raise RuntimeError("❌ Error: ktc.compile() did not return a valid .nef file path!")
|
|
|
|
# 確保 .nef 檔案存在
|
|
if not os.path.exists(nef_model_path):
|
|
raise RuntimeError(f"❌ Error: NEF model was not generated at {nef_model_path}! Please check your compilation process.")
|
|
|
|
# 確保 NEF 檔案儲存到指定目錄
|
|
nef_save_path = os.path.join(onnx_dir, os.path.basename(nef_model_path))
|
|
if os.path.exists(nef_model_path):
|
|
shutil.copy(nef_model_path, nef_save_path)
|
|
else:
|
|
raise RuntimeError(f"❌ Error: NEF model was expected at {nef_model_path}, but it does not exist!") # 同樣使用 shutil.move
|
|
|
|
if not os.path.exists(nef_save_path):
|
|
raise RuntimeError(f"\u274c Error: NEF model was not generated! Please check your compilation process.")
|
|
|
|
print("\n\u2705 Compile done! NEF file saved to:", nef_save_path)
|