增加向量模型转化成安卓可运行模型的程序

faiss_vector_db/build_faiss_index.py

@@ -267,8 +267,11 @@ class FAISSIndexBuilder:
         
         # 保存标签映射
         labels_path = os.path.join(output_dir, 'labels.pkl')
+        labels_json_path = os.path.join(output_dir, 'labels.json')
         with open(labels_path, 'wb') as f:
             pickle.dump(self.labels, f)
+        with open(labels_json_path, "w", encoding="utf-8") as f:
+            json.dump(self.labels, f, ensure_ascii=False)
         print(f"标签映射已保存到: {labels_path}")
         
         # 保存类别信息
@@ -288,8 +291,16 @@ class FAISSIndexBuilder:
         
         # 保存特征向量（可选，用于调试）
         embeddings_path = os.path.join(output_dir, 'embeddings.npy')
+        embeddings_json_path = os.path.join(output_dir, 'embeddings.json')
         embeddings_array = np.array(self.embeddings)
         np.save(embeddings_path, embeddings_array)
+
+        # 转成 Python 列表（JSON 不支持 numpy.array 直接序列化）
+        embeddings_list = embeddings_array.tolist()
+        # 保存成 JSON 文件
+        with open(embeddings_json_path, "w", encoding="utf-8") as f:
+            json.dump(embeddings_list, f, ensure_ascii=False)
+
         print(f"特征向量已保存到: {embeddings_path}")
     
     def build_complete_index(self, train_dir: str, output_dir: str, 

toAndroid/toAndroidEmbedding.py

@@ -0,0 +1,196 @@
+import torch
+import torch.nn.functional as F
+import os
+import sys
+
+# 添加项目根目录到路径
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+sys.path.insert(0, parent_dir)
+
+from net.resnet_embedding import create_resnet50_embedding, create_mobile_resnet50_embedding
+
+
+def main():
+    print("开始转换ResNet50 Embedding模型为移动端格式...")
+    
+    # 1. 加载训练好的embedding模型权重
+    base_model = create_resnet50_embedding(embedding_dim=512, pretrained=True)
+    model_path = "../model/embedding_20250917_145342/best_embedding_model.pth"
+    
+    if not os.path.exists(model_path):
+        print(f"错误：模型文件不存在 {model_path}")
+        return
+    
+    # 加载模型权重
+    checkpoint = torch.load(model_path, map_location='cpu')
+    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+        # 如果保存的是完整的checkpoint
+        base_model.load_state_dict(checkpoint['model_state_dict'])
+        print("✓ 从checkpoint加载模型权重成功")
+    else:
+        # 如果保存的是纯模型权重
+        base_model.load_state_dict(checkpoint)
+        print("✓ 基础模型权重加载成功")
+    
+    # 设置为评估模式
+    base_model.eval()
+    
+    # 2. 创建移动端模型并复制权重
+    # 直接使用基础模型，但在forward中包含预处理
+    mobile_model = base_model  # 使用同一个模型实例
+    print("✓ 移动端模型创建成功")
+    
+    # 3. 测试模型
+    # 创建示例输入 (模拟Android端输入：2x3x224x224，值范围0-1，使用batch_size=2避免BatchNorm问题)
+    example_input = torch.rand(2, 3, 224, 224, dtype=torch.float32)
+    
+    with torch.no_grad():
+        # 测试基础模型
+        base_output = base_model(example_input)
+        print(f"✓ 基础模型测试成功，输出形状: {base_output.shape}")
+        print(f"  基础模型输出值范围: [{base_output.min():.3f}, {base_output.max():.3f}]")
+        print(f"  基础模型输出L2范数: {torch.norm(base_output, p=2, dim=1)}")
+        
+        # 测试移动端模型（使用forward_mobile方法）
+        mobile_output = mobile_model.forward_mobile(example_input)
+        print(f"✓ 移动端模型测试成功，输出形状: {mobile_output.shape}")
+        print(f"  移动端模型输出值范围: [{mobile_output.min():.3f}, {mobile_output.max():.3f}]")
+        print(f"  移动端模型输出L2范数: {torch.norm(mobile_output, p=2, dim=1)}")
+        
+        # 检查输出是否一致（由于预处理不同，可能会有差异）
+        print("✓ 移动端模型测试完成（包含预处理）")
+    
+    # 4. 创建专门的移动端模型类用于TorchScript转换
+    class MobileEmbeddingModel(torch.nn.Module):
+        def __init__(self, base_model):
+            super().__init__()
+            self.base_model = base_model
+        
+        def forward(self, x):
+            # 在forward中包含预处理
+            return self.base_model.forward_mobile(x)
+    
+    # 创建移动端包装模型
+    mobile_wrapper = MobileEmbeddingModel(base_model)
+    mobile_wrapper.eval()
+    
+    # 5. 转换为 TorchScript
+    try:
+        # 为TorchScript转换创建单个样本输入 (1x3x224x224)
+        single_input = torch.rand(1, 3, 224, 224, dtype=torch.float32)
+        
+        # 使用trace方法转换
+        traced_model = torch.jit.trace(mobile_wrapper, single_input)
+        
+        # 保存模型
+        output_path = "../model/embedding_20250917_145342/best_embedding_model_mobile.pt"
+        traced_model.save(output_path)
+        print(f"✓ TorchScript模型保存成功: {output_path}")
+        
+        # 验证保存的模型
+        loaded_model = torch.jit.load(output_path)
+        with torch.no_grad():
+            # 使用单个样本测试保存的模型
+            loaded_output = loaded_model(single_input)
+            print(f"✓ 保存的模型验证成功，输出形状: {loaded_output.shape}")
+            print(f"  保存模型输出L2范数: {torch.norm(loaded_output, p=2, dim=1)}")
+            
+            # 用移动端模型处理单个样本进行比较
+            mobile_single_output = mobile_wrapper(single_input)
+            
+            # 检查输出是否一致
+            if torch.allclose(mobile_single_output, loaded_output, atol=1e-6):
+                print("✓ 模型输出一致性验证通过")
+            else:
+                print("⚠ 警告：模型输出存在差异")
+                print(f"  差异最大值: {torch.max(torch.abs(mobile_single_output - loaded_output))}")
+        
+        # 测试不同尺寸的输入
+        print("\n测试不同输入尺寸...")
+        test_sizes = [(1, 3, 256, 256), (1, 3, 320, 320), (1, 3, 128, 128)]
+        
+        for size in test_sizes:
+            test_input = torch.rand(*size, dtype=torch.float32)
+            try:
+                with torch.no_grad():
+                    test_output = loaded_model(test_input)
+                    print(f"✓ 输入尺寸 {size} 测试成功，输出形状: {test_output.shape}")
+                    print(f"  输出L2范数: {torch.norm(test_output, p=2, dim=1)}")
+            except Exception as e:
+                print(f"✗ 输入尺寸 {size} 测试失败: {e}")
+        
+    except Exception as e:
+        print(f"✗ TorchScript转换失败: {e}")
+        import traceback
+        traceback.print_exc()
+        return
+    
+    # 5. 模型信息统计
+    print("\n" + "="*60)
+    print("模型信息统计:")
+    print("="*60)
+    
+    # 计算模型大小
+    model_size = os.path.getsize(output_path) / (1024 * 1024)  # MB
+    print(f"模型文件大小: {model_size:.2f} MB")
+    
+    # 计算参数数量
+    total_params = sum(p.numel() for p in base_model.parameters())
+    trainable_params = sum(p.numel() for p in base_model.parameters() if p.requires_grad)
+    print(f"总参数数量: {total_params:,}")
+    print(f"可训练参数数量: {trainable_params:,}")
+    
+    print("\n" + "="*60)
+    print("移动端Embedding模型转换完成！")
+    print("="*60)
+    print("使用说明：")
+    print("1. Android端输入图片应为 [batch, 3, height, width] 格式")
+    print("2. 像素值范围：0-1 (float32)")
+    print("3. 模型会自动处理：")
+    print("   - 缩放到 224x224")
+    print("   - ImageNet标准化 (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])")
+    print("4. 输出：512维L2归一化的embedding向量")
+    print("5. 可用于：")
+    print("   - 图片相似度计算 (余弦相似度)")
+    print("   - 图片检索和匹配")
+    print("   - 特征向量数据库构建")
+    print("="*60)
+    
+    # 6. 生成示例代码
+    print("\nAndroid端使用示例 (Java/Kotlin):")
+    print("```java")
+    print("// 加载模型")
+    print("Module module = LiteModuleLoader.load(assetFilePath(\"best_embedding_model_mobile.pt\"));")
+    print("")
+    print("// 准备输入 (假设bitmap已转换为float数组，值范围0-1)")
+    print("float[] inputArray = ...; // [1, 3, height, width]")
+    print("Tensor inputTensor = Tensor.fromBlob(inputArray, new long[]{1, 3, height, width});")
+    print("")
+    print("// 推理")
+    print("IValue[] outputs = module.forward(IValue.from(inputTensor)).toTuple();")
+    print("Tensor outputTensor = outputs[0].toTensor();")
+    print("float[] embedding = outputTensor.getDataAsFloatArray(); // 512维向量")
+    print("")
+    print("// 计算相似度")
+    print("float similarity = cosineSimilarity(embedding1, embedding2);")
+    print("```")
+
+
+def test_model_conversion():
+    """
+    测试模型转换的完整流程
+    """
+    print("开始模型转换测试...")
+    
+    try:
+        main()
+        print("✓ 模型转换测试完成")
+    except Exception as e:
+        print(f"✗ 模型转换测试失败: {e}")
+        import traceback
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    main()