@tensorflow/tfjs-core
Version:
Hardware-accelerated JavaScript library for machine intelligence
54 lines • 2.18 kB
JavaScript
/**
* @license
* Copyright 2019 Google LLC. All Rights Reserved.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
* =============================================================================
*/
import { concat } from '../concat';
import { mul } from '../mul';
import { op } from '../operation';
import { enclosingPowerOfTwo } from '../signal_ops_util';
import { slice } from '../slice';
import { rfft } from '../spectral/rfft';
import { frame } from './frame';
import { hannWindow } from './hann_window';
/**
* Computes the Short-time Fourier Transform of signals
* See: https://en.wikipedia.org/wiki/Short-time_Fourier_transform
*
* ```js
* const input = tf.tensor1d([1, 1, 1, 1, 1])
* tf.signal.stft(input, 3, 1).print();
* ```
* @param signal 1-dimensional real value tensor.
* @param frameLength The window length of samples.
* @param frameStep The number of samples to step.
* @param fftLength The size of the FFT to apply.
* @param windowFn A callable that takes a window length and returns 1-d tensor.
*
* @doc {heading: 'Operations', subheading: 'Signal', namespace: 'signal'}
*/
function stft_(signal, frameLength, frameStep, fftLength, windowFn = hannWindow) {
if (fftLength == null) {
fftLength = enclosingPowerOfTwo(frameLength);
}
const framedSignal = frame(signal, frameLength, frameStep);
const windowedSignal = mul(framedSignal, windowFn(frameLength));
const output = [];
for (let i = 0; i < framedSignal.shape[0]; i++) {
output.push(rfft(slice(windowedSignal, [i, 0], [1, frameLength]), fftLength));
}
return concat(output);
}
export const stft = op({ stft_ });
//# sourceMappingURL=stft.js.map