From 422556c4a3317bf6cbc9680ac7809152093a4de7 Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Fri, 30 Jun 2023 08:31:37 -0700
Subject: [PATCH] Internal change

PiperOrigin-RevId: 544663494
---
 .../audio/spectrogram_calculator.cc           | 45 ++++++++++---------
 .../audio/spectrogram_calculator.proto        |  2 +-
 2 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/mediapipe/calculators/audio/spectrogram_calculator.cc b/mediapipe/calculators/audio/spectrogram_calculator.cc
index fbdbbab0a..7f6528ec1 100644
--- a/mediapipe/calculators/audio/spectrogram_calculator.cc
+++ b/mediapipe/calculators/audio/spectrogram_calculator.cc
@@ -210,6 +210,23 @@ REGISTER_CALCULATOR(SpectrogramCalculator);
 // Factor to convert ln(SQUARED_MAGNITUDE) to deciBels = 10.0/ln(10.0).
 const float SpectrogramCalculator::kLnSquaredMagnitudeToDb = 4.342944819032518;
 
+namespace {
+std::unique_ptr<audio_dsp::WindowFunction> MakeWindowFun(
+    const SpectrogramCalculatorOptions::WindowType window_type) {
+  switch (window_type) {
+    // The cosine window and square root of Hann are equivalent.
+    case SpectrogramCalculatorOptions::COSINE:
+    case SpectrogramCalculatorOptions::SQRT_HANN:
+      return std::make_unique<audio_dsp::CosineWindow>();
+    case SpectrogramCalculatorOptions::HANN:
+      return std::make_unique<audio_dsp::HannWindow>();
+    case SpectrogramCalculatorOptions::HAMMING:
+      return std::make_unique<audio_dsp::HammingWindow>();
+  }
+  return nullptr;
+}
+}  // namespace
+
 absl::Status SpectrogramCalculator::Open(CalculatorContext* cc) {
   SpectrogramCalculatorOptions spectrogram_options =
       cc->Options<SpectrogramCalculatorOptions>();
@@ -266,28 +283,14 @@ absl::Status SpectrogramCalculator::Open(CalculatorContext* cc) {
 
   output_scale_ = spectrogram_options.output_scale();
 
-  std::vector<double> window;
-  switch (spectrogram_options.window_type()) {
-    case SpectrogramCalculatorOptions::COSINE:
-      audio_dsp::CosineWindow().GetPeriodicSamples(frame_duration_samples_,
-                                                   &window);
-      break;
-    case SpectrogramCalculatorOptions::HANN:
-      audio_dsp::HannWindow().GetPeriodicSamples(frame_duration_samples_,
-                                                 &window);
-      break;
-    case SpectrogramCalculatorOptions::HAMMING:
-      audio_dsp::HammingWindow().GetPeriodicSamples(frame_duration_samples_,
-                                                    &window);
-      break;
-    case SpectrogramCalculatorOptions::SQRT_HANN: {
-      audio_dsp::HannWindow().GetPeriodicSamples(frame_duration_samples_,
-                                                 &window);
-      absl::c_transform(window, window.begin(),
-                        [](double x) { return std::sqrt(x); });
-      break;
-    }
+  auto window_fun = MakeWindowFun(spectrogram_options.window_type());
+  if (window_fun == nullptr) {
+    return absl::Status(absl::StatusCode::kInvalidArgument,
+                        absl::StrCat("Invalid window type ",
+                                     spectrogram_options.window_type()));
   }
+  std::vector<double> window;
+  window_fun->GetPeriodicSamples(frame_duration_samples_, &window);
 
   // Propagate settings down to the actual Spectrogram object.
   spectrogram_generators_.clear();
diff --git a/mediapipe/calculators/audio/spectrogram_calculator.proto b/mediapipe/calculators/audio/spectrogram_calculator.proto
index ddfca1d1c..d8bca3f76 100644
--- a/mediapipe/calculators/audio/spectrogram_calculator.proto
+++ b/mediapipe/calculators/audio/spectrogram_calculator.proto
@@ -68,7 +68,7 @@ message SpectrogramCalculatorOptions {
     HANN = 0;
     HAMMING = 1;
     COSINE = 2;
-    SQRT_HANN = 4;
+    SQRT_HANN = 4;  // Alias of COSINE.
   }
   optional WindowType window_type = 6 [default = HANN];