#!/bin/bash
# SCRFD PTH -> ONNX -> RKNN 转换脚本

cd /home/orangepi/apps/OrangePi3588Media/utils

echo "=== 步骤1: 导出 ONNX ==="
python3 << 'PYEOF'
import torch
import sys
sys.path.insert(0, '/home/orangepi/apps/OrangePi3588Media')

# 加载 SCRFD 模型
print("加载 SCRFD 模型...")
from insightface.model_zoo.scrfd import SCRFD
import cv2
import numpy as np

# 创建模型实例
model = SCRFD(model_file='scrfd_2.5g_kps.pth')
model.eval()

# 检查是否有 CUDA
if torch.cuda.is_available():
    model = model.cuda()
    device = 'cuda'
else:
    device = 'cpu'
    
print(f"使用设备: {device}")

# 创建 dummy input (640x640)
dummy_input = torch.randn(1, 3, 640, 640)
if device == 'cuda':
    dummy_input = dummy_input.cuda()

# 导出 ONNX
print("导出 ONNX...")
torch.onnx.export(
    model,
    dummy_input,
    'scrfd_2.5g_kps.onnx',
    input_names=['input'],
    output_names=['score_8', 'score_16', 'score_32', 'bbox_8', 'bbox_16', 'bbox_32', 'kps_8', 'kps_16', 'kps_32'],
    dynamic_axes={'input': {0: 'batch', 2: 'height', 3: 'width'}},
    opset_version=11,
    do_constant_folding=True
)
print('ONNX 导出成功!')
PYEOF

echo ""
echo "=== 步骤2: 简化 ONNX ==="
python3 -m onnxsim scrfd_2.5g_kps.onnx scrfd_2.5g_kps_sim.onnx

echo ""
echo "=== 步骤3: 转换为 RKNN ==="
cd /home/orangepi/apps/OrangePi3588Media/models

python3 << 'PYEOF'
from rknn.api import RKNN

print("初始化 RKNN...")
rknn = RKNN(verbose=False)

# 配置 - SCRFD 使用标准预处理
rknn.config(
    target_platform='rk3588',
    mean_values=[[127.5, 127.5, 127.5]],  # 归一化到 -1~1
    std_values=[[128.0, 128.0, 128.0]],
    quantized_dtype='w8a8',
)

print("加载简化后的 ONNX...")
rknn.load_onnx(model='../utils/scrfd_2.5g_kps_sim.onnx')

print("构建 RKNN 模型...")
rknn.build(do_quantization=True, dataset='./dataset.txt')

print("导出 RKNN...")
rknn.export_rknn('scrfd_2.5g_kps.rknn')
print("SCRFD 640 RKNN 模型转换完成!")
rknn.release()
PYEOF

echo ""
echo "=== 完成 ==="
ls -lh /home/orangepi/apps/OrangePi3588Media/models/scrfd_2.5g_kps.rknn