Home | History | Annotate | Download | only in kernels
      1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
      2 
      3 Licensed under the Apache License, Version 2.0 (the "License");
      4 you may not use this file except in compliance with the License.
      5 You may obtain a copy of the License at
      6 
      7     http://www.apache.org/licenses/LICENSE-2.0
      8 
      9 Unless required by applicable law or agreed to in writing, software
     10 distributed under the License is distributed on an "AS IS" BASIS,
     11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 See the License for the specific language governing permissions and
     13 limitations under the License.
     14 ==============================================================================*/
     15 
     16 // Class for generating spectrogram slices from a waveform.
     17 // Initialize() should be called before calls to other functions.  Once
     18 // Initialize() has been called and returned true, The Compute*() functions can
     19 // be called repeatedly with sequential input data (ie. the first element of the
     20 // next input vector directly follows the last element of the previous input
     21 // vector). Whenever enough audio samples are buffered to produce a
     22 // new frame, it will be placed in output. Output is cleared on each
     23 // call to Compute*(). This class is thread-unsafe, and should only be
     24 // called from one thread at a time.
     25 // With the default parameters, the output of this class should be very
     26 // close to the results of the following MATLAB code:
     27 // overlap_samples = window_length_samples - step_samples;
     28 // window = hann(window_length_samples, 'periodic');
     29 // S = abs(spectrogram(audio, window, overlap_samples)).^2;
     30 
     31 #ifndef TENSORFLOW_CORE_KERNELS_SPECTROGRAM_H_
     32 #define TENSORFLOW_CORE_KERNELS_SPECTROGRAM_H_
     33 
     34 #include <complex>
     35 #include <deque>
     36 #include <vector>
     37 
     38 #include "third_party/fft2d/fft.h"
     39 #include "tensorflow/core/framework/op_kernel.h"
     40 #include "tensorflow/core/framework/tensor.h"
     41 
     42 namespace tensorflow {
     43 
     44 class Spectrogram {
     45  public:
     46   Spectrogram() : initialized_(false) {}
     47   ~Spectrogram() {}
     48 
     49   // Initializes the class with a given window length and step length
     50   // (both in samples). Internally a Hann window is used as the window
     51   // function. Returns true on success, after which calls to Process()
     52   // are possible. window_length must be greater than 1 and step
     53   // length must be greater than 0.
     54   bool Initialize(int window_length, int step_length);
     55 
     56   // Initialize with an explicit window instead of a length.
     57   bool Initialize(const std::vector<double>& window, int step_length);
     58 
     59   // Processes an arbitrary amount of audio data (contained in input)
     60   // to yield complex spectrogram frames. After a successful call to
     61   // Initialize(), Process() may be called repeatedly with new input data
     62   // each time.  The audio input is buffered internally, and the output
     63   // vector is populated with as many temporally-ordered spectral slices
     64   // as it is possible to generate from the input.  The output is cleared
     65   // on each call before the new frames (if any) are added.
     66   //
     67   // The template parameters can be float or double.
     68   template <class InputSample, class OutputSample>
     69   bool ComputeComplexSpectrogram(
     70       const std::vector<InputSample>& input,
     71       std::vector<std::vector<std::complex<OutputSample>>>* output);
     72 
     73   // This function works as the one above, but returns the power
     74   // (the L2 norm, or the squared magnitude) of each complex value.
     75   template <class InputSample, class OutputSample>
     76   bool ComputeSquaredMagnitudeSpectrogram(
     77       const std::vector<InputSample>& input,
     78       std::vector<std::vector<OutputSample>>* output);
     79 
     80   // Return reference to the window function used internally.
     81   const std::vector<double>& GetWindow() const { return window_; }
     82 
     83   // Return the number of frequency channels in the spectrogram.
     84   int output_frequency_channels() const { return output_frequency_channels_; }
     85 
     86  private:
     87   template <class InputSample>
     88   bool GetNextWindowOfSamples(const std::vector<InputSample>& input,
     89                               int* input_start);
     90   void ProcessCoreFFT();
     91 
     92   int fft_length_;
     93   int output_frequency_channels_;
     94   int window_length_;
     95   int step_length_;
     96   bool initialized_;
     97   int samples_to_next_step_;
     98 
     99   std::vector<double> window_;
    100   std::vector<double> fft_input_output_;
    101   std::deque<double> input_queue_;
    102 
    103   // Working data areas for the FFT routines.
    104   std::vector<int> fft_integer_working_area_;
    105   std::vector<double> fft_double_working_area_;
    106 
    107   TF_DISALLOW_COPY_AND_ASSIGN(Spectrogram);
    108 };
    109 
    110 }  // namespace tensorflow
    111 
    112 #endif  // TENSORFLOW_CORE_KERNELS_SPECTROGRAM_H_
    113