STDC/tools/kneron/onnx2nef720.py

import ktc
import numpy as np
import os
import onnx
import shutil
from PIL import Image

# === 1. 設定路徑與參數 ===
onnx_dir = 'work_dirs/meconfig8/'  # 你的 onnx存放路徑
onnx_path = os.path.join(onnx_dir, 'latest.onnx')
data_path = "data724362"  # 測試圖片資料夾
imgsz_w, imgsz_h = 724, 362  # 輸入圖片尺寸，跟ONNX模型要求一致

# === 2. 建立輸出資料夾 ===
os.makedirs(onnx_dir, exist_ok=True)

# === 3. 載入並優化 ONNX 模型 ===
print("🔄 Loading and optimizing 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)

# === 4. 檢查 ONNX 輸入尺寸是否符合要求 ===
input_tensor = m.graph.input[0]
input_shape = [dim.dim_value for dim in input_tensor.type.tensor_type.shape.dim]
print(f"📏 ONNX Input Shape: {input_shape}")

expected_shape = [1, 3, imgsz_h, imgsz_w]  # (N, C, H, W)

if input_shape != expected_shape:
    raise ValueError(f"❌ Error: ONNX input shape {input_shape} does not match expected {expected_shape}.")

# === 5. 設定 Kneron 模型編譯參數 ===
print("📐 Configuring model for KL720...")
km = ktc.ModelConfig(20008, "0001", "720", onnx_model=m)

# （可選）模型效能評估
eval_result = km.evaluate()
print("\n📊 NPU Performance Evaluation:\n" + str(eval_result))

# === 6. 準備圖片資料 ===
print("🖼️ Preparing image data...")
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"❌ No images found in {data_path}!")

print(f"✅ Found {len(files_found)} images in {data_path}")

input_name = input_tensor.name
img_list = []

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")):
            continue
        try:
            img = Image.open(fullpath).convert("RGB")
            img = Image.fromarray(np.array(img)[..., ::-1])  # RGB ➔ BGR
            img_np = np.array(img.resize((imgsz_w, imgsz_h), Image.BILINEAR)).astype(np.float32).copy()
            img_np = img_np / 256.0 - 0.5
            img_np = np.transpose(img_np, (2, 0, 1))  # HWC ➔ CHW
            img_np = np.expand_dims(img_np, axis=0)   # CHW ➔ NCHW
            img_list.append(img_np)
            print(f"✅ Processed: {fullpath}")
        except Exception as e:
            print(f"❌ Failed to process {fullpath}: {e}")

if not img_list:
    raise RuntimeError("❌ Error: No valid images were processed!")

# === 7. BIE 量化分析 ===
print("📦 Running fixed-point analysis (BIE)...")
bie_model_path = km.analysis({input_name: img_list})
bie_save_path = os.path.join(onnx_dir, os.path.basename(bie_model_path))
shutil.copy(bie_model_path, bie_save_path)

if not os.path.exists(bie_save_path):
    raise RuntimeError("❌ Error: BIE model was not generated!")

print("✅ BIE model saved to:", bie_save_path)

# === 8. 編譯 NEF 模型 ===
print("⚙️ Compiling NEF model...")
nef_model_path = ktc.compile([km])
nef_save_path = os.path.join(onnx_dir, os.path.basename(nef_model_path))
shutil.copy(nef_model_path, nef_save_path)

if not os.path.exists(nef_save_path):
    raise RuntimeError("❌ Error: NEF model was not generated!")

print("✅ NEF compile done!")
print("📁 NEF file saved to:", nef_save_path)