67 lines
2.2 KiB
Python
67 lines
2.2 KiB
Python
import os
|
|
import numpy as np
|
|
import torch
|
|
import ktc # Kneron Toolchain
|
|
import onnx
|
|
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 = "/workspace/yolov5/data50"
|
|
|
|
# 確保 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 模型的輸入名稱**
|
|
onnx_model_path = "/workspace/yolov5/runs/train/exp24/weights/latest.opt.onnx"
|
|
m = onnx.load(onnx_model_path)
|
|
input_name = m.graph.input[0].name # 確保 key 與 ONNX input name 一致
|
|
|
|
km = ktc.ModelConfig(20008, "0001", "720", onnx_model=onnx_model_path)
|
|
# 存儲預處理後的圖片數據
|
|
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)
|