Function spectrogram

tools/audio/source/audio.cpp:385–455

}

VARP spectrogram(VARP waveform, const SpectrogramParams *params) {
    int pad_left = 0, pad_right = 0, pad_mode = REFLECT;
    int n_fft = 400, hop_length = 0, win_length = 0, window_type = HANNING;
    bool center = false, normalized = false;
    float power = 2.0;
    if (params) {
        pad_left    = params->pad_left;
        pad_right   = params->pad_right;
        center      = params->center;
        pad_mode    = params->pad_mode;
        n_fft       = params->n_fft;
        hop_length  = params->hop_length;
        win_length  = params->win_length;
        window_type = params->window_type;
        normalized  = params->normalized;
        power       = params->power;
    }
    if (pad_left > 1 || pad_right > 1) {
        waveform = MNN::Express::_Pad(waveform, _var<int>({pad_left, pad_right}, {2}), MNN::Express::CONSTANT);
    }
    if (center) {
        waveform = MNN::Express::_Pad(waveform, _var<int>({n_fft / 2, n_fft / 2}, {2}), static_cast<MNN::Express::PadValueMode>(pad_mode));
    }
    waveform = _Reshape(waveform, {1, -1, 1});
    hop_length = hop_length ? hop_length : n_fft / 2;
    win_length = win_length ? win_length : n_fft;
    VARP window;
    switch (window_type) {
        case HANNING:
            window = hann_window(win_length);
            break;
        case HAMMING:
            window = hamming_window(win_length);
            break;
        default:
            window = hann_window(win_length);
            break;
    }
    std::unique_ptr<OpT> op(new OpT);
    op->type      = OpType_Stft;
    op->main.type = OpParameter_StftParam;
    auto param = new StftParamT;
    param->abs = true;
    op->main.value = param;
    EXPRP stftexpr = Expr::create(std::move(op), {waveform, _Scalar<int>(hop_length), window});
    int frame_size = win_length > 0 ? win_length : n_fft;
    int nstfts = ((waveform->getInfo()->dim[1] - frame_size) / hop_length) + 1;
    int dft_unique_bins = n_fft / 2 + 1;
    auto specgram = MNN::Express::Variable::create(stftexpr);
    specgram = _Square(specgram);
    auto startsDims = std::vector<int>{0, 0, 0, 0};
    auto starts1Dims = std::vector<int>{0, 0, 0, 1};
    auto sizeDims = std::vector<int>{1, nstfts, dft_unique_bins, 1};
    auto startVar = _Const(startsDims.data(), {4}, NCHW, halide_type_of<int>());
    auto start1Var = _Const(starts1Dims.data(), {4}, NCHW, halide_type_of<int>());
    auto sizeVar = _Const(sizeDims.data(), {4}, NCHW, halide_type_of<int>());
    auto specgramReal = _Slice(specgram, startVar, sizeVar);
    auto specgramVirt = _Slice(specgram, start1Var, sizeVar);