Begin to add C++ runtime for 3d-speaker models

k2-fsa · Jan 10, 2024 · 70fad62 · 70fad62
1 parent 7faeb24
commit 70fad62
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 15 deletions.
diff --git a/...aker-embedding-extractor-wespeaker-impl.h → ...peaker-embedding-extractor-general-impl.h b/...aker-embedding-extractor-wespeaker-impl.h → ...peaker-embedding-extractor-general-impl.h
@@ -1,23 +1,23 @@
-// sherpa-onnx/csrc/speaker-embedding-extractor-wespeaker-impl.h
+// sherpa-onnx/csrc/speaker-embedding-extractor-general-impl.h
 //
 // Copyright (c)  2023  Xiaomi Corporation
 
-#ifndef SHERPA_ONNX_CSRC_SPEAKER_EMBEDDING_EXTRACTOR_WESPEAKER_IMPL_H_
-#define SHERPA_ONNX_CSRC_SPEAKER_EMBEDDING_EXTRACTOR_WESPEAKER_IMPL_H_
+#ifndef SHERPA_ONNX_CSRC_SPEAKER_EMBEDDING_EXTRACTOR_GENERAL_IMPL_H_
+#define SHERPA_ONNX_CSRC_SPEAKER_EMBEDDING_EXTRACTOR_GENERAL_IMPL_H_
 #include <algorithm>
 #include <memory>
 #include <utility>
 #include <vector>
 
 #include "sherpa-onnx/csrc/speaker-embedding-extractor-impl.h"
-#include "sherpa-onnx/csrc/speaker-embedding-extractor-wespeaker-model.h"
+#include "sherpa-onnx/csrc/speaker-embedding-extractor-model.h"
 
 namespace sherpa_onnx {
 
-class SpeakerEmbeddingExtractorWeSpeakerImpl
+class SpeakerEmbeddingExtractorGeneralImpl
     : public SpeakerEmbeddingExtractorImpl {
  public:
-  explicit SpeakerEmbeddingExtractorWeSpeakerImpl(
+  explicit SpeakerEmbeddingExtractorGeneralImpl(
       const SpeakerEmbeddingExtractorConfig &config)
       : model_(config) {}
 
@@ -71,9 +71,9 @@ class SpeakerEmbeddingExtractorWeSpeakerImpl
   }
 
  private:
-  SpeakerEmbeddingExtractorWeSpeakerModel model_;
+  SpeakerEmbeddingExtractorModel model_;
 };
 
 }  // namespace sherpa_onnx
 
-#endif  // SHERPA_ONNX_CSRC_SPEAKER_EMBEDDING_EXTRACTOR_WESPEAKER_IMPL_H_
+#endif  // SHERPA_ONNX_CSRC_SPEAKER_EMBEDDING_EXTRACTOR_GENERAL_IMPL_H_
diff --git a/sherpa-onnx/csrc/speaker-embedding-extractor-impl.cc b/sherpa-onnx/csrc/speaker-embedding-extractor-impl.cc
@@ -5,14 +5,15 @@
 
 #include "sherpa-onnx/csrc/macros.h"
 #include "sherpa-onnx/csrc/onnx-utils.h"
-#include "sherpa-onnx/csrc/speaker-embedding-extractor-wespeaker-impl.h"
+#include "sherpa-onnx/csrc/speaker-embedding-extractor-general-impl.h"
 
 namespace sherpa_onnx {
 
 namespace {
 
 enum class ModelType {
   kWeSpeaker,
+  k3dSpeaker,
   kUnkown,
 };
 
@@ -49,6 +50,8 @@ static ModelType GetModelType(char *model_data, size_t model_data_length,
 
   if (model_type.get() == std::string("wespeaker")) {
     return ModelType::kWeSpeaker;
+  } else if (model_type.get() == std::string("3d-speaker")) {
+    return ModelType::k3dSpeaker;
   } else {
     SHERPA_ONNX_LOGE("Unsupported model_type: %s", model_type.get());
     return ModelType::kUnkown;
@@ -68,7 +71,9 @@ SpeakerEmbeddingExtractorImpl::Create(
 
   switch (model_type) {
     case ModelType::kWeSpeaker:
-      return std::make_unique<SpeakerEmbeddingExtractorWeSpeakerImpl>(config);
+      // fall through
+    case ModelType::k3dSpeaker:
+      return std::make_unique<SpeakerEmbeddingExtractorGeneralImpl>(config);
     case ModelType::kUnkown:
       SHERPA_ONNX_LOGE(
           "Unknown model type in for speaker embedding extractor!");

diff --git a/sherpa-onnx/csrc/speaker-embedding-extractor-model.cc b/sherpa-onnx/csrc/speaker-embedding-extractor-model.cc
@@ -1,6 +1,6 @@
 // sherpa-onnx/csrc/speaker-embedding-extractor-model.cc
 //
-// Copyright (c)  2023  Xiaomi Corporation
+// Copyright (c)  2023-2024  Xiaomi Corporation
 
 #include "sherpa-onnx/csrc/speaker-embedding-extractor-model.h"
 
@@ -11,7 +11,7 @@
 #include "sherpa-onnx/csrc/macros.h"
 #include "sherpa-onnx/csrc/onnx-utils.h"
 #include "sherpa-onnx/csrc/session.h"
-#include "sherpa-onnx/csrc/speaker-embedding-extractor-model-metadata.h"
+#include "sherpa-onnx/csrc/speaker-embedding-extractor-model-meta-data.h"
 
 namespace sherpa_onnx {
 
@@ -70,7 +70,7 @@ class SpeakerEmbeddingExtractorModel::Impl {
 
     std::string framework;
     SHERPA_ONNX_READ_META_DATA_STR(framework, "framework");
-    if (framework != "wespeaker" || framework != "3d-speaker") {
+    if (framework != "wespeaker" && framework != "3d-speaker") {
       SHERPA_ONNX_LOGE("Expect a wespeaker or a 3d-speaker model, given: %s",
                        framework.c_str());
       exit(-1);

diff --git a/sherpa-onnx/csrc/speaker-embedding-extractor-model.h b/sherpa-onnx/csrc/speaker-embedding-extractor-model.h
@@ -1,13 +1,13 @@
 // sherpa-onnx/csrc/speaker-embedding-extractor-model.h
 //
-// Copyright (c)  2023  Xiaomi Corporation
+// Copyright (c)  2023-2024  Xiaomi Corporation
 #ifndef SHERPA_ONNX_CSRC_SPEAKER_EMBEDDING_EXTRACTOR_MODEL_H_
 #define SHERPA_ONNX_CSRC_SPEAKER_EMBEDDING_EXTRACTOR_MODEL_H_
 
 #include <memory>
 
 #include "onnxruntime_cxx_api.h"  // NOLINT
-#include "sherpa-onnx/csrc/speaker-embedding-extractor-model-metadata.h"
+#include "sherpa-onnx/csrc/speaker-embedding-extractor-model-meta-data.h"
 #include "sherpa-onnx/csrc/speaker-embedding-extractor.h"
 
 namespace sherpa_onnx {