Create infer_contrast.py

infer_contrast.py

@@ -0,0 +1,51 @@
+import argparse
+import functools
+
+import numpy as np
+import torch
+
+from utils.reader import load_audio
+from utils.utility import add_arguments, print_arguments
+
+parser = argparse.ArgumentParser(description=__doc__)
+add_arg = functools.partial(add_arguments, argparser=parser)
+add_arg('threshold',        float,   0.71,                    '判断是否为同一个人的阈值')
+add_arg('input_shape',      str,    '(1, 257, 257)',          '数据输入的形状')
+add_arg('model_path',       str,    'models_large/resnet34.pth',    '预测模型的路径')
+# args = parser.parse_args()
+args =parser.parse_known_args()[0]
+
+print_arguments(args)
+
+device = torch.device("cuda")
+
+# 加载模型
+# model = torch.jit.load(args.model_path)
+model = torch.load(args.model_path)
+# model.to(device)
+model.eval()
+
+
+# 预测音频
+def infer(audio_path):
+    input_shape = eval(args.input_shape)
+    data = load_audio(audio_path, mode='infer', spec_len=input_shape[2])
+    data = data[np.newaxis, :]
+    data = torch.tensor(data, dtype=torch.float32)
+    # 执行预测
+    feature = model(data)
+    return feature.data.cpu().numpy()
+
+
+def run(audio1,audio2):
+    # 要预测的两个人的音频文件
+    feature1 = infer(audio1)[0]
+    feature2 = infer(audio2)[0]
+    # 对角余弦值
+    dist = np.dot(feature1, feature2) / (np.linalg.norm(feature1) * np.linalg.norm(feature2))
+    if dist > args.threshold:
+        result = "Speaker1 和 Speaker2 为同一个人，相似度为：%f" % (dist)
+    else:
+        result = "Speaker1 和 Speaker2 为同一个人，相似度为：%f" % (dist)
+        
+    return result