88 lines
2.0 KiB
Bash
Executable File
88 lines
2.0 KiB
Bash
Executable File
#!/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
|