Use TFLite as default for Separator

jinay1991 · Nov 1, 2020 · 39f7cba · 39f7cba
1 parent 66c218a
commit 39f7cba
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/spleeter/argument_parser/cli_options.h b/spleeter/argument_parser/cli_options.h
@@ -33,7 +33,7 @@ struct CLIOptions
 
     /// @brief Inference Engine Parameters (contains model_path, input/output tensor names)
     InferenceEngineParameters inference_engine_params{
-        "external/models/5stems/saved_model",
+        "external/models/5stems.tflite",
         "waveform",
         {"strided_slice_18", "strided_slice_38", "strided_slice_48", "strided_slice_28", "strided_slice_58"},
         "spleeter:5stems"};

diff --git a/spleeter/separator.cpp b/spleeter/separator.cpp
@@ -45,7 +45,7 @@ Separator::Separator(const InferenceEngineParameters& inference_engine_param, co
       inference_engine_{},
       waveform_name_{internal::GetWaveformNames(inference_engine_param.configuration)}
 {
-    inference_engine_.SelectInferenceEngine(InferenceEngineType::kTensorFlow, inference_engine_param);
+    inference_engine_.SelectInferenceEngine(InferenceEngineType::kTensorFlowLite, inference_engine_param);
     inference_engine_.Init();
 }
 

diff --git a/spleeter/test/separator_tests.cpp b/spleeter/test/separator_tests.cpp
@@ -27,7 +27,7 @@ class SeparatorTest : public ::testing::TestWithParam<std::int32_t>
     {
         const auto stem{std::to_string(GetParam()) + "stems"};
         cli_options_.configuration = "spleeter:" + stem;
-        cli_options_.inference_engine_params.model_path = "external/models/" + stem + "/saved_model";
+        cli_options_.inference_engine_params.model_path = "external/models/" + stem + "/" + stem + ".tflite";
         cli_options_.inference_engine_params.output_tensor_names = GetOutputTensorNames(cli_options_.configuration);
         cli_options_.inference_engine_params.configuration = cli_options_.configuration;
 
@@ -36,7 +36,7 @@ class SeparatorTest : public ::testing::TestWithParam<std::int32_t>
         unit_ = std::make_unique<Separator>(cli_options_.inference_engine_params, cli_options_.mwf);
     }
 
-    std::vector<std::string> GetOutputTensorNames(const std::string& configuration)
+    static std::vector<std::string> GetOutputTensorNames(const std::string& configuration)
     {
         auto output_tensor_names = std::vector<std::string>{};
         if (configuration == "spleeter:2stems")