Yolov5s/kneron/ktc720.py
2026-03-11 16:13:59 +08:00

72 lines
2.4 KiB
Python

import ktc
import numpy as np
import os
import onnx
from PIL import Image
import torch
from yolov5_preprocess import Yolov5_preprocess
import kneron_preprocessing
onnx_path = 'runs/train/exp24/weights/best_simplified.onnx'
m = onnx.load(onnx_path)
m = ktc.onnx_optimizer.onnx2onnx_flow(m)
onnx.save(m,'latest.opt.onnx')
km = ktc.ModelConfig(20008, "0001", "720", 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 = "data50"
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"❌ Error: No images found in {data_path}! Please check your dataset.")
print(f"✅ 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"⚠️ 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"✅ Processed: {fullpath}")
img_list.append(img_data)
except Exception as e:
print(f"❌ Failed to process {fullpath}: {e}")
# **確保 img_list 不是空的**
if not img_list:
raise ValueError("❌ Error: No valid images were processed! Please check the image paths and formats.")
# **執行 BIE 量化**
bie_model_path = km.analysis({input_name: img_list})
# **確認 BIE 模型是否生成**
if not os.path.exists(bie_model_path):
raise RuntimeError(f"❌ Error: BIE model was not generated! Please check your quantization process.")
# 顯示成功訊息
print("\n✅ Fixed-point analysis done! BIE model saved to:", bie_model_path)
# 確保 `km` 已經初始化,並且 `.bie` 模型已生成
nef_model_path = ktc.compile([km])
# 顯示成功訊息
print("\n✅ Compile done! NEF file saved to:", nef_model_path)