From e2dedcbfe569d4a33ad24ac77fee51a2ed53d5b2 Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Wed, 18 Jan 2023 19:40:19 -0800
Subject: [PATCH] Add SQRT_HANN window type to both SpectrogramCalculator and
 InverseSpectrogramCalculator.

PiperOrigin-RevId: 503041493
---
 mediapipe/calculators/audio/spectrogram_calculator.cc    | 7 +++++++
 mediapipe/calculators/audio/spectrogram_calculator.proto | 1 +
 2 files changed, 8 insertions(+)

diff --git a/mediapipe/calculators/audio/spectrogram_calculator.cc b/mediapipe/calculators/audio/spectrogram_calculator.cc
index c038c0cd7..bd4d8f3bf 100644
--- a/mediapipe/calculators/audio/spectrogram_calculator.cc
+++ b/mediapipe/calculators/audio/spectrogram_calculator.cc
@@ -280,6 +280,13 @@ absl::Status SpectrogramCalculator::Open(CalculatorContext* cc) {
       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;
+    }
   }
 
   // Propagate settings down to the actual Spectrogram object.
diff --git a/mediapipe/calculators/audio/spectrogram_calculator.proto b/mediapipe/calculators/audio/spectrogram_calculator.proto
index 8e1e18051..ddfca1d1c 100644
--- a/mediapipe/calculators/audio/spectrogram_calculator.proto
+++ b/mediapipe/calculators/audio/spectrogram_calculator.proto
@@ -68,6 +68,7 @@ message SpectrogramCalculatorOptions {
     HANN = 0;
     HAMMING = 1;
     COSINE = 2;
+    SQRT_HANN = 4;
   }
   optional WindowType window_type = 6 [default = HANN];