audio_effects_demo_directory.cpp

audio_effects_demo_directory.cpp

/*###############################################################################
#
# Copyright 2020 NVIDIA Corporation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
# the Software, and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
###############################################################################*/

#include <cstring>
#include <iostream>
#include <iomanip>
#include <memory>
#include <sstream>
#include <chrono>
#include <thread>
#include <string>

#include <utils/wave_reader/waveReadWrite.hpp>
#include <utils/config_reader/ConfigReader.hpp>

#include <nvAudioEffects.h>
#include <windows.h>

namespace {

    bool ReadWavFile(const std::string &filename, uint32_t expected_sample_rate, std::vector<float> *data,
                     int align_samples) {
        CWaveFileRead wave_file(filename);
        if (wave_file.isValid() == false) {
            std::cout << "wave_file.isValid() == false " << std::endl;
            return false;
        }
        const float *raw_data_array = wave_file.GetFloatPCMData();
        std::cout << "Total number of samples: " << wave_file.GetNumSamples() << std::endl;
        std::cout << "Size in bytes: " << wave_file.GetRawPCMDataSizeInBytes() << std::endl;
        std::cout << "Sample rate: " << wave_file.GetSampleRate() << std::endl;

        auto bits_per_sample = wave_file.GetBitsPerSample();
        std::cout << "Bits/sample: " << bits_per_sample << std::endl;

        if (wave_file.GetSampleRate() != expected_sample_rate) {
            std::cout << "Sample rate mismatch" << std::endl;
            return false;
        }
        if (wave_file.GetWaveFormat().nChannels != 1) {
            std::cout << "Channel count needs to be 1" << std::endl;
            return false;
        }

        if (align_samples != -1) {
            int num_frames = wave_file.GetNumSamples() / align_samples;
            if (wave_file.GetNumSamples() % align_samples) {
                num_frames++;
            }

            // allocate potentially a bigger sized buffer to align it to requested
            data->resize(num_frames * align_samples);
        } else {
            data->resize(wave_file.GetNumSamples(), 0.f);
        }
        std::copy(raw_data_array, raw_data_array + wave_file.GetNumSamples(), data->data());
        return true;
    }

    /* allowed sample rates */
    const std::vector<uint32_t> kAllowedSampleRates = {16000, 48000};
} // namespace

class EffectsDemoApp {
public:
    bool run(const std::string &input_wav);

private:
    // Sample rate config
    uint32_t sample_rate_ = 48000;
    // Intensity_ratio config
    float intensity_ratio_ = 1.0f;
    // inited from configuration
    bool real_time_ = false;
};

bool EffectsDemoApp::run(const std::string &input_wav) {
    if (real_time_) {
        std::cout << "App will run in real time mode ..." << std::endl;
    }

    int num_effects;
    NvAFX_EffectSelector *effects;
    if (NvAFX_GetEffectList(&num_effects, &effects) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "NvAFX_GetEffectList() failed" << std::endl;
        return false;
    }

    std::cout << "Total Effects supported: " << num_effects << std::endl;
    for (int i = 0; i < num_effects; ++i) {
        std::cout << "(" << i + 1 << ") " << effects[i] << std::endl;
    }

    NvAFX_Handle handle;
    std::string effect = "dereverb_denoiser";
    if (strcmp(effect.c_str(), "denoiser") == 0) {
        if (NvAFX_CreateEffect(NVAFX_EFFECT_DENOISER, &handle) != NVAFX_STATUS_SUCCESS) {
            std::cerr << "NvAFX_CreateEffect() failed" << std::endl;
            return false;
        }
    } else if (strcmp(effect.c_str(), "dereverb") == 0) {
        if (NvAFX_CreateEffect(NVAFX_EFFECT_DEREVERB, &handle) != NVAFX_STATUS_SUCCESS) {
            std::cerr << "NvAFX_CreateEffect() failed" << std::endl;
            return false;
        }
    } else if (strcmp(effect.c_str(), "dereverb_denoiser") == 0) {
        if (NvAFX_CreateEffect(NVAFX_EFFECT_DEREVERB_DENOISER, &handle) != NVAFX_STATUS_SUCCESS) {
            std::cerr << "NvAFX_CreateEffect() failed" << std::endl;
            return false;
        }
    } else {
        std::cerr << "NvAFX_CreateEffect() failed. Invalid Effect Value : " << effect << std::endl;
        return false;
    }

    // If the system has multiple supported GPUs, then the application can either
    // use CUDA driver APIs or CUDA runtime APIs to enumerate the GPUs and select one based on the application's requirements
    // or offload the responsibility to SDK to select the GPU by setting NVAFX_PARAM_USE_DEFAULT_GPU as 1
    /*if (NvAFX_SetU32(handle, NVAFX_PARAM_USE_DEFAULT_GPU, 1) != NVAFX_STATUS_SUCCESS) {
      std::cerr << "NvAFX_SetBool(NVAFX_PARAM_USE_DEFAULT_GPU " << ") failed" << std::endl;
      return false;
    }*/

    if (NvAFX_SetU32(handle, NVAFX_PARAM_SAMPLE_RATE, sample_rate_) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "NvAFX_SetU32(Sample Rate: " << sample_rate_ << ") failed" << std::endl;
        return false;
    }

    if (NvAFX_SetFloat(handle, NVAFX_PARAM_INTENSITY_RATIO, intensity_ratio_) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "NvAFX_SetFloat(Intensity Ratio: " << intensity_ratio_ << ") failed" << std::endl;
        return false;
    }
    std::string model_file = "models/ampere/dereverb_denoiser_48k.trtpkg";

    if (NvAFX_SetString(handle, NVAFX_PARAM_MODEL_PATH, model_file.c_str()) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "NvAFX_SetString() failed" << std::endl;
        return false;
    }

    // Another option could be to use cudaGetDeviceCount for num
    int num_supported_devices = 0;
    if (NvAFX_GetSupportedDevices(handle, &num_supported_devices, nullptr) != NVAFX_STATUS_OUTPUT_BUFFER_TOO_SMALL) {
        std::cerr << "Could not get number of supported devices" << std::endl;
        return false;
    }

    std::cout << "Number of supported devices for this model: " << num_supported_devices << std::endl;

    std::vector<int> ret(num_supported_devices);
    if (NvAFX_GetSupportedDevices(handle, &num_supported_devices, ret.data()) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "No supported devices found" << std::endl;
        return false;
    }

    std::cout << "Devices supported (sorted by preference)" << std::endl;
    for (int device : ret) {
        std::cout << "- " << device << std::endl;
    }

    std::vector<float> audio_data;
    std::cout << "Loading effect"
              << " ... ";
    if (NvAFX_Load(handle) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "NvAFX_Load() failed" << std::endl;
        return false;
    }
    std::cout << "Done" << std::endl;
    unsigned num_channels, num_samples_per_frame;
    if (NvAFX_GetU32(handle, NVAFX_PARAM_NUM_CHANNELS, &num_channels) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "NvAFX_GetU32() failed" << std::endl;
        return false;
    }
    if (NvAFX_GetU32(handle, NVAFX_PARAM_NUM_SAMPLES_PER_FRAME, &num_samples_per_frame) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "NvAFX_GetU32() failed" << std::endl;
        return false;
    }
    float intensity_ratio_local;
    if (NvAFX_GetFloat(handle, NVAFX_PARAM_INTENSITY_RATIO, &intensity_ratio_local) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "NvAFX_GetFloat() failed" << std::endl;
        return false;
    }

    std::cout << "  Effect properties : " << std::endl
              << "  Channels            : " << num_channels << std::endl
              << "  Samples per frame   : " << num_samples_per_frame << std::endl
              << "  Intensity Ratio   : " << intensity_ratio_local << std::endl;

    if (!ReadWavFile(input_wav, sample_rate_, &audio_data, num_samples_per_frame)) {
        std::cerr << "Unable to read wav file: " << input_wav << std::endl;
        return false;
    }
    std::cout << "Input wav file: " << input_wav << std::endl
              << "Total " << audio_data.size() << " samples read" << std::endl;

    // std::string output_wav = config_reader.GetConfigValue(kConfigFileOutputVariable);
    std::string output_wav = input_wav + ".out.wav";
    // std::string output_wav = "hogeout.wav";

    CWaveFileWrite wav_write(output_wav, sample_rate_, num_channels, 32, true);
    float frame_in_secs = static_cast<float>(num_samples_per_frame) / static_cast<float>(sample_rate_);
    float total_run_time = 0.f;
    float total_audio_duration = 0.f;
    float checkpoint = 0.1f;
    float expected_audio_duration = static_cast<float>(audio_data.size()) / static_cast<float>(sample_rate_);
    auto frame = std::make_unique<float[]>(num_samples_per_frame);

    std::string progress_bar = "[          ] ";
    std::cout << "Processed: " << progress_bar << "0%\r";
    std::cout.flush();

    // wav data is already padded to align to num_samples_per_frame by ReadWavFile()
    for (size_t offset = 0; offset < audio_data.size(); offset += num_samples_per_frame) {
        const float *input[1];
        float *output[1];
        input[0] = &audio_data.data()[offset];
        output[0] = frame.get();

        auto start_tick = std::chrono::high_resolution_clock::now();
        if (NvAFX_Run(handle, input, output, num_samples_per_frame, num_channels) != NVAFX_STATUS_SUCCESS) {
            std::cerr << "NvAFX_Run() failed" << std::endl;
            return false;
        }

        auto run_end_tick = std::chrono::high_resolution_clock::now();
        total_run_time += (std::chrono::duration<float>(run_end_tick - start_tick)).count();
        total_audio_duration += frame_in_secs;

        if ((total_audio_duration / expected_audio_duration) >= checkpoint) {
            progress_bar[checkpoint * 10] = '=';
            std::cout << "Processed: " << progress_bar << checkpoint * 100.f << "%" << (checkpoint >= 1 ? "\n" : "\r");
            std::cout.flush();
            checkpoint += 0.1f;
        }

        wav_write.writeChunk(frame.get(), num_samples_per_frame * sizeof(float));

        if (real_time_) {
            auto end_tick = std::chrono::high_resolution_clock::now();
            std::chrono::duration<float> elapsed = end_tick - start_tick;
            float sleep_time_secs = frame_in_secs - elapsed.count();
            std::this_thread::sleep_for(std::chrono::milliseconds(static_cast<int>(sleep_time_secs * 1000)));
        }
    }

    std::cout << "Processing time " << std::setprecision(2) << total_run_time
              << " secs for " << total_audio_duration << std::setprecision(2)
              << " secs audio file (" << total_run_time / total_audio_duration
              << " secs processing time per sec of audio)" << std::endl;

    if (real_time_) {
        std::cout << "Note: App ran in real time mode i.e. simulated the input data rate of a mic" << std::endl
                  << "'Processing time' could be less then actual run time" << std::endl;
    }

    wav_write.commitFile();
    std::cout << "Output wav file written. " << output_wav << std::endl
              << "Total " << audio_data.size() << " samples written" << std::endl;

    if (NvAFX_DestroyEffect(handle) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "NvAFX_Release() failed" << std::endl;
        return false;
    }

    return true;
}

void ShowHelpAndExit(const char *bad_option) {
    std::ostringstream oss;
    if (bad_option) {
        oss << "Error parsing \"" << bad_option << "\"" << std::endl;
    }
    std::cout << "Command Line Options:" << std::endl
              << "-c Config file" << std::endl;
}

#ifdef _MSC_VER
#define strncasecmp _strnicmp
#define strcasecmp _stricmp
#endif

void ParseCommandLine(int argc, char *argv[], std::string *config_file) {
    if (argc == 1) {
        ShowHelpAndExit(nullptr);
    }

    for (int i = 1; i < argc; i++) {
        if (!strcasecmp(argv[i], "-h")) {
            ShowHelpAndExit(nullptr);
        }
        if (!strcasecmp(argv[i], "-c")) {
            if (++i == argc || !config_file->empty()) {
                ShowHelpAndExit("-f");
            }
            config_file->assign(argv[i]);
            continue;
        }

        ShowHelpAndExit(argv[i]);
    }
}

int main(int argc, char *argv[]) {

    EffectsDemoApp app;

    WIN32_FIND_DATA data;
    HANDLE hFind = FindFirstFile("C:/Program Files/NVIDIA Corporation/NVIDIA Audio Effects SDK/noisy-data/*",
                                 &data); // DIRECTORY

    if (hFind != INVALID_HANDLE_VALUE) {
        do {
            std::cout << "input: " << data.cFileName << std::endl;

            std::string fn = std::string(data.cFileName);
            if (fn.length() > 5) {

                std::cout << "fn: " << fn << std::endl;

                if (app.run("C:/Program Files/NVIDIA Corporation/NVIDIA Audio Effects SDK/noisy-data/" + fn)) {
                    std::cout << fn + " finiehd!!" << std::endl;
                } else {
                    std::cout << fn + " error..." << std::endl;
                }
            }

        } while (FindNextFile(hFind, &data));
        FindClose(hFind);
    }
}