import argparse
import os
import shutil
import time

import numpy as np

from utils import model, utils
from utils.record import RecordAudio

parser = argparse.ArgumentParser()
parser.add_argument('--audio_db',    default='audio_db/',              type=str,   help='音频库的路径')
parser.add_argument('--threshold',   default=0.7,                      type=float, help='判断是否为同一个人的阈值')
parser.add_argument('--model_path',  default=r'models/resnet34-56.h5', type=str,   help='模型的路径')
args = parser.parse_args()

person_feature = []
person_name = []

# 获取模型
network_eval = model.vggvox_resnet2d_icassp(input_dim=(257, None, 1), mode='eval')

# 加载预训练模型
network_eval.load_weights(os.path.join(args.model_path), by_name=True)
print('==> successfully loading model {}.'.format(args.model_path))


# 预测获取声纹特征
def predict(path):
    specs = utils.load_data(path, mode='eval')
    specs = np.expand_dims(np.expand_dims(specs, 0), -1)
    feature = network_eval.predict(specs)[0]
    return feature


# 加载要识别的音频库
def load_audio_db(audio_db_path):
    start = time.time()
    audios = os.listdir(audio_db_path)
    for audio in audios:
        path = os.path.join(audio_db_path, audio)
        name = audio[:-4]
        feature = predict(path)
        person_name.append(name)
        person_feature.append(feature)
        print("Loaded %s audio." % name)
    end = time.time()
    print('加载音频库完成，消耗时间：%fms' % (round((end - start) * 1000)))


# 识别声纹
def recognition(path):
    name = ''
    pro = 0
    feature = predict(path)
    for i, person_f in enumerate(person_feature):
        # 计算相识度
        dist = np.dot(feature, person_f.T)
        if dist > pro:
            pro = dist
            name = person_name[i]
    return name, pro


# 声纹注册
def register(path, user_name):
    save_path = os.path.join(args.audio_db, user_name + os.path.basename(path)[-4:])
    shutil.move(path, save_path)
    feature = predict(save_path)
    person_name.append(user_name)
    person_feature.append(feature)


if __name__ == '__main__':
    load_audio_db(args.audio_db)
    record_audio = RecordAudio()

    while True:
        select_fun = int(input("请选择功能，0为注册音频到声纹库，1为执行声纹识别："))
        if select_fun == 0:
            audio_path = record_audio.record()
            name = input("请输入该音频用户的名称：")
            if name == '': continue
            register(audio_path, name)
        elif select_fun == 1:
            audio_path = record_audio.record()
            name, p = recognition(audio_path)
            if p > args.threshold:
                print("识别说话的为：%s，相似度为：%f" % (name, p))
            else:
                print("音频库没有该用户的语音")
        else:
            print('请正确选择功能')

[(11, 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), False, 'import argparse\n'), (21, 'utils.model.vggvox_resnet2d_icassp', 'model.vggvox_resnet2d_icassp', ([], {'input_dim': '(257, None, 1)', 'mode': '"""eval"""'}), False, 'from utils import model, utils\n'), (24, 'os.path.join', 'os.path.join', (['args.model_path'], {}), False, 'import os\n'), (30, 'utils.utils.load_data', 'utils.load_data', (['path'], {'mode': '"""eval"""'}), False, 'from utils import model, utils\n'), (38, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (39, 'os.listdir', 'os.listdir', (['audio_db_path'], {}), False, 'import os\n'), (47, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (68, 'shutil.move', 'shutil.move', (['path', 'save_path'], {}), False, 'import shutil\n'), (76, 'utils.record.RecordAudio', 'RecordAudio', ([], {}), False, 'from utils.record import RecordAudio\n'), (31, 'numpy.expand_dims', 'np.expand_dims', (['specs', '(0)'], {}), True, 'import numpy as np\n'), (41, 'os.path.join', 'os.path.join', (['audio_db_path', 'audio'], {}), False, 'import os\n'), (58, 'numpy.dot', 'np.dot', (['feature', 'person_f.T'], {}), True, 'import numpy as np\n'), (67, 'os.path.basename', 'os.path.basename', (['path'], {}), False, 'import os\n')]