change model with a sc model

convert_model.py

@@ -1,134 +1,62 @@
 #!/usr/bin/env python3
 """
-ArcFace ONNX to TFLite 转换脚本
+ONNX 人脸模型转 TFLite 脚本（基于 onnx2tf）
 
-使用方法:
-    1. 先安装依赖：pip install onnx onnxruntime tensorflow
-    2. 运行转换：python convert_model.py
+环境要求（推荐使用已有 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 numpy as np
+import shutil
 import subprocess
 import sys
+import tempfile
 
-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"
+def main() -> int:
+    if len(sys.argv) != 3:
+        print(__doc__)
+        return 1
 
-    print("下载 ArcFace 模型 (IResNet-50, Glint360K 数据集)...")
-    print(f"目标路径：{output_path}")
+    onnx_path = sys.argv[1]
+    tflite_path = sys.argv[2]
 
-    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}")
+    if not os.path.isfile(onnx_path):
+        print(f"错误：找不到输入模型 {onnx_path}")
+        return 1
 
+    work_dir = tempfile.mkdtemp(prefix="onnx2tf_")
     try:
-        import onnx
-        from onnx_tf.backend import prepare
+        print(f"转换 {onnx_path} -> {tflite_path}")
+        subprocess.check_call(
+            [sys.executable, "-m", "onnx2tf", "-i", onnx_path, "-o", work_dir]
+        )
 
-        # 加载 ONNX 模型
-        onnx_model = onnx.load(onnx_path)
-        onnx.checker.check_model(onnx_model)
+        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])
 
-        # 打印模型信息
-        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]}")
+        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)
 
-        # 转换为 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()
+    sys.exit(main())