添加新的人脸识别模型

convert_model.py

@@ -0,0 +1,134 @@
+#!/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()