From e87bd0c3fb53afece68554f28c415653f0311e52 Mon Sep 17 00:00:00 2001
From: Binbin Zhang <binbzha@qq.com>
Date: Wed, 8 Nov 2023 09:56:18 +0800
Subject: [PATCH] [cli] save embedding to file (#218)

---
 README.md                | 2 +-
 wespeaker/cli/speaker.py | 9 ++++++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index c379ac40..52ae32cc 100644
--- a/README.md
+++ b/README.md
@@ -23,7 +23,7 @@ pip install git+https://github.com/wenet-e2e/wespeaker.git
 **Command-line usage** (use `-h` for parameters):
 
 ``` sh
-$ wespeaker --task embedding --audio_file audio.wav
+$ wespeaker --task embedding --audio_file audio.wav --output_file embedding.txt
 $ wespeaker --task similarity --audio_file audio.wav --audio_file2 audio2.wav
 $ wespeaker --task diarization --audio_file audio.wav # TODO
 ```
diff --git a/wespeaker/cli/speaker.py b/wespeaker/cli/speaker.py
index 6dfaeca9..308fef50 100644
--- a/wespeaker/cli/speaker.py
+++ b/wespeaker/cli/speaker.py
@@ -129,6 +129,8 @@ def get_args():
     parser.add_argument('--vad',
                         action='store_true',
                         help='whether to do VAD or not')
+    parser.add_argument('--output_file',
+                        help='output file to save speaker embedding')
     args = parser.parse_args()
     return args
 
@@ -137,7 +139,12 @@ def main():
     args = get_args()
     model = load_model(args.language, args.resample_rate)
     if args.task == 'embedding':
-        print(model.extract_embedding(args.audio_file, args.vad))
+        embedding = model.extract_embedding(args.audio_file, args.vad)
+        if embedding is not None:
+            np.savetxt(args.output_file, embedding)
+            print('Succeed, see {}'.format(args.output_file))
+        else:
+            print('Fails to extract embedding')
     elif args.task == 'similarity':
         print(model.compute_similarity(args.audio_file, args.audio_file2))
     elif args.task == 'diarization':