63 lines
1.8 KiB
Python
63 lines
1.8 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
ONNX 人脸模型转 TFLite 脚本（基于 onnx2tf）
|
||
|
||
环境要求（推荐使用已有 onnx2tf 的 conda 环境）：
|
||
conda activate base # 已安装 onnx2tf + ai-edge-litert
|
||
或 pip install onnx2tf ai-edge-litert onnx onnxruntime
|
||
|
||
用法:
|
||
python convert_model.py <input.onnx> <output.tflite>
|
||
|
||
示例:
|
||
python convert_model.py backup_models/buffalo_sc/w600k_mbf.onnx app/app/src/main/assets/mobilefacenet.tflite
|
||
"""
|
||
|
||
import os
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
|
||
|
||
def main() -> int:
|
||
if len(sys.argv) != 3:
|
||
print(__doc__)
|
||
return 1
|
||
|
||
onnx_path = sys.argv[1]
|
||
tflite_path = sys.argv[2]
|
||
|
||
if not os.path.isfile(onnx_path):
|
||
print(f"错误：找不到输入模型 {onnx_path}")
|
||
return 1
|
||
|
||
work_dir = tempfile.mkdtemp(prefix="onnx2tf_")
|
||
try:
|
||
print(f"转换 {onnx_path} -> {tflite_path}")
|
||
subprocess.check_call(
|
||
[sys.executable, "-m", "onnx2tf", "-i", onnx_path, "-o", work_dir]
|
||
)
|
||
|
||
base = os.path.splitext(os.path.basename(onnx_path))[0]
|
||
converted = os.path.join(work_dir, f"{base}_float32.tflite")
|
||
if not os.path.isfile(converted):
|
||
candidates = [
|
||
f for f in os.listdir(work_dir) if f.endswith("_float32.tflite")
|
||
]
|
||
if not candidates:
|
||
print("错误：onnx2tf 未生成 float32 tflite，请查看上方日志")
|
||
return 1
|
||
converted = os.path.join(work_dir, candidates[0])
|
||
|
||
os.makedirs(os.path.dirname(tflite_path) or ".", exist_ok=True)
|
||
shutil.copyfile(converted, tflite_path)
|
||
print(f"完成：{tflite_path}（{os.path.getsize(tflite_path) / 1024 / 1024:.2f} MB）")
|
||
return 0
|
||
finally:
|
||
shutil.rmtree(work_dir, ignore_errors=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|