135 lines
4.2 KiB
Python
135 lines
4.2 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
ArcFace ONNX to TFLite 转换脚本
|
||
|
||
使用方法:
|
||
1. 先安装依赖：pip install onnx onnxruntime tensorflow
|
||
2. 运行转换：python convert_model.py
|
||
"""
|
||
|
||
import os
|
||
import numpy as np
|
||
import subprocess
|
||
import sys
|
||
|
||
def install_requirements():
|
||
"""安装必要的包"""
|
||
packages = ['onnx', 'onnxruntime', 'tensorflow', 'tf2onnx', 'onnx-tf']
|
||
print("检查并安装依赖包...")
|
||
for pkg in packages:
|
||
try:
|
||
__import__(pkg.replace('-', '_'))
|
||
print(f" ✓ {pkg} 已安装")
|
||
except ImportError:
|
||
print(f" 安装 {pkg}...")
|
||
subprocess.check_call([sys.executable, '-m', 'pip', 'install', pkg, '-q'])
|
||
print("依赖包安装完成！\n")
|
||
|
||
def download_model():
|
||
"""下载 ArcFace 模型"""
|
||
model_url = "https://drive.google.com/uc?export=download&id=1gnt6P3jaiwfevV4hreWHPu0Mive5VRyP"
|
||
output_path = "/Users/liushuming/projects/app/app/src/main/assets/arcface_ir50_glint360k.onnx"
|
||
|
||
print("下载 ArcFace 模型 (IResNet-50, Glint360K 数据集)...")
|
||
print(f"目标路径：{output_path}")
|
||
|
||
import urllib.request
|
||
urllib.request.urlretrieve(model_url, output_path)
|
||
|
||
size_mb = os.path.getsize(output_path) / 1024 / 1024
|
||
print(f"下载完成！模型大小：{size_mb:.2f} MB\n")
|
||
|
||
return output_path
|
||
|
||
def convert_onnx_to_tflite(onnx_path, output_path):
|
||
"""将 ONNX 模型转换为 TFLite"""
|
||
print(f"转换 ONNX -> TFLite...")
|
||
print(f"输入：{onnx_path}")
|
||
print(f"输出：{output_path}")
|
||
|
||
try:
|
||
import onnx
|
||
from onnx_tf.backend import prepare
|
||
|
||
# 加载 ONNX 模型
|
||
onnx_model = onnx.load(onnx_path)
|
||
onnx.checker.check_model(onnx_model)
|
||
|
||
# 打印模型信息
|
||
inputs = onnx_model.graph.input
|
||
outputs = onnx_model.graph.output
|
||
print(f"\n模型输入：{inputs[0].name}, 形状：{[d.dim_value for d in inputs[0].type.tensor_type.shape.dim]}")
|
||
print(f"模型输出：{outputs[0].name}, 形状：{[d.dim_value for d in outputs[0].type.tensor_type.shape.dim]}")
|
||
|
||
# 转换为 TensorFlow
|
||
print("\n转换为 TensorFlow 格式...")
|
||
tf_backend = prepare(onnx_model)
|
||
tf_graph = tf_backend.tf_graph
|
||
|
||
# 保存为 SavedModel
|
||
import tensorflow as tf
|
||
saved_model_dir = output_path.replace('.tflite', '_saved_model')
|
||
|
||
# 清理已存在的目录
|
||
import shutil
|
||
if os.path.exists(saved_model_dir):
|
||
shutil.rmtree(saved_model_dir)
|
||
|
||
# 保存 graph
|
||
with tf.Graph().as_default() as graph:
|
||
tf.import_graph_def(tf_graph, name="")
|
||
tf.saved_model.save(tf.keras.models.Model(inputs=graph.get_tensor_by_name('input:0'),
|
||
outputs=graph.get_tensor_by_name('output:0')),
|
||
saved_model_dir)
|
||
|
||
# 转换为 TFLite
|
||
print("转换为 TFLite 格式...")
|
||
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
|
||
converter.optimizations = [tf.lite.Optimize.DEFAULT]
|
||
|
||
tflite_model = converter.convert()
|
||
|
||
with open(output_path, 'wb') as f:
|
||
f.write(tflite_model)
|
||
|
||
# 清理临时文件
|
||
shutil.rmtree(saved_model_dir)
|
||
|
||
size_mb = os.path.getsize(output_path) / 1024 / 1024
|
||
print(f"\n✓ 转换完成！")
|
||
print(f" TFLite 模型：{output_path}")
|
||
print(f" 模型大小：{size_mb:.2f} MB")
|
||
|
||
return True
|
||
|
||
except ImportError as e:
|
||
print(f"\n错误：缺少依赖包 - {e}")
|
||
print("请运行：pip install onnx-tf tensorflow")
|
||
return False
|
||
except Exception as e:
|
||
print(f"\n转换错误：{e}")
|
||
return False
|
||
|
||
def main():
|
||
print("=" * 60)
|
||
print("ArcFace 模型下载和转换工具")
|
||
print("=" * 60)
|
||
|
||
# 安装依赖
|
||
install_requirements()
|
||
|
||
# 下载模型
|
||
onnx_path = download_model()
|
||
|
||
# 转换为 TFLite
|
||
output_path = "/Users/liushuming/projects/app/app/src/main/assets/facenet.tflite"
|
||
convert_onnx_to_tflite(onnx_path, output_path)
|
||
|
||
print("\n" + "=" * 60)
|
||
print("完成！")
|
||
print("模型已保存到：app/src/main/assets/facenet.tflite")
|
||
print("=" * 60)
|
||
|
||
if __name__ == "__main__":
|
||
main()
|