AudioDataset.cpp

Go to the documentation of this file.

#include <chrono>
 
namespace krotos
{
 
AudioDataset::AudioDataset() : juce::Thread("AudioDatasetThread")
{
    // sort here as using std::binary_search later
    std::sort(CatIDSet.begin(), CatIDSet.end());
 
    startThread(); // set a priority?
}
 
AudioDataset::~AudioDataset() { stopThread(4000); }
 
std::vector<File> AudioDataset::findValidFiles() const
{
    std::vector<File> paths;
 
    AudioFormatManager formatManager;
    formatManager.registerBasicFormats();
 
    const auto root = getFactoryAssetDirectory();
    for (const auto& path : File(root).findChildFiles(2, true, "*.wav", juce::File::FollowSymlinks::no))
    {
        // extract UCS CatID from filename
        const auto filename = path.getFileName().toStdString();
        const auto key = extractCatID(filename);
 
        // check if we have a valid CatID tag
        const auto success = std::binary_search(CatIDSet.begin(), CatIDSet.end(), key);
        if (success)
        {
            paths.push_back(path);
        }
    }
    return paths;
}
 
AudioBuffer<float> AudioDataset::resampleAudioBuffer(const AudioBuffer<float>& buffer, double sampleRate,
                                                     double targetRate)
{
    if ((int)sampleRate == (int)targetRate)
        return buffer;
 
    const double ratio = sampleRate / targetRate;
    AudioBuffer<float> resampledBuffer;
    resampledBuffer.setSize(1, (int)(buffer.getNumSamples() / ratio));
    auto inputData = buffer.getReadPointer(0);
    auto outputData = resampledBuffer.getWritePointer(0);
    LagrangeInterpolator resampler;
    resampler.process(ratio, inputData, outputData, resampledBuffer.getNumSamples());
    return resampledBuffer;
}
 
void AudioDataset::run()
{
    if (runAnalysis())
    {
        const auto tick = std::chrono::steady_clock::now();
 
        AudioFormatManager formatManager;
        formatManager.registerBasicFormats();
        AudioEmbedding model;
 
        const auto paths = findValidFiles();
        const auto file_count = paths.size();
 
        for (const auto& path : paths)
        {
            // extract UCS CatID from filename
            const auto filename = path.getFileName().toStdString();
            const auto key = extractCatID(filename);
 
            std::unique_ptr<juce::AudioFormatReader> reader(formatManager.createReaderFor(path));
            if (reader.get() != nullptr)
            {
                // load up to 3 seconds of audio
                const auto numSamples =
                        std::min(static_cast<int>(reader->lengthInSamples), static_cast<int>(3 * reader->sampleRate));
                const auto duration = static_cast<float>(reader->lengthInSamples) / reader->sampleRate;
                int64 readerStartSample = 0;
 
                // To-Do: investigate suitable duration threshold
                if (duration > 10.f)
                {
                    // analyse the start of short Foley files,
                    // but shift to centre of long files to avoid fade in
                    readerStartSample = static_cast<int64>(reader->lengthInSamples / 2) - (numSamples / 2);
                }
 
                // opportunity to exit before audio analysis
                if (threadShouldExit())
                    break;
 
                // read audio segment into buffer
                juce::AudioBuffer<float> buffer(1, numSamples);
                reader->read(&buffer, 0, numSamples, readerStartSample, true, false);
 
                // resample to target samplerate of 48 kHz
                buffer = resampleAudioBuffer(buffer, reader->sampleRate, 48000.0);
 
                // extract the audio embedding
                const auto embedding = model.forward(buffer);
 
                // store the path and embedding together
                const auto data = std::make_pair(path.getFullPathName().toStdString(), embedding);
                m_dataset[key].push_back(data);
            }
 
            // another chance to exit
            if (threadShouldExit())
                break;
        }
 
        const auto tock = std::chrono::steady_clock::now();
        const auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(tock - tick).count();
        DBG("analysis time: " << elapsed << " ms");
        DBG("file count: " << file_count);
 
        const auto success = writeCache();
        if (success)
        {
            const auto outfile = getDatasetFile().getFullPathName().toStdString();
            DBG("analysis cached to: " << outfile);
        }
        else
        {
            DBG("caching analysis failed");
        }
    }
}
 
bool AudioDataset::assetsChanged()
{
    const auto loaded_paths = getPathsSet();
 
    const auto paths = findValidFiles();
    std::set<std::string> available_paths;
    for (const auto& path : paths)
        available_paths.insert(path.getFullPathName().toStdString());
 
    return (available_paths != loaded_paths);
}
 
bool AudioDataset::runAnalysis()
{
    if (cacheExists())
    {
        DBG("loading previous analysis from cache");
        bool success = readCache();
        if (success)
        {
            // check if factory assets changed
            if (assetsChanged())
            {
                // factory assets changed since we last cached analysis
                // rerun the analysis
                m_dataset.clear();
                success = deleteCache();
                if (success)
                {
                    DBG("rerunning analysis...");
                    return true;
                }
                else
                {
                    // uh-oh
                    return false;
                }
            }
        }
        else
        {
            // cache not loaded
            // rerun the analysis
            m_dataset.clear();
            success = deleteCache();
            if (success)
            {
                DBG("rerunning analysis...");
                return true;
            }
            else
            {
                // uh-oh
                return false;
            }
        }
    }
    else
    {
        DBG("running analysis...");
        return true;
    }
 
    return false;
}
 
std::set<std::string> AudioDataset::getPathsSet() const
{
    std::set<std::string> paths;
    for (const auto& key : m_dataset)
    {
        for (const auto& element : key.second)
        {
            const auto path = element.first;
            paths.insert(path);
        }
    }
    return paths;
}
 
std::string AudioDataset::fixKeyErrors(std::string key) const
{
    // workarounds for some of the factory assets
    if (key == "FOLEYFeet")
    {
        key = "FOLYFeet";
    }
    else if (key.find("CMPTKey") != std::string::npos)
    {
        key = "CMPTKey";
    }
    return key;
}
 
bool AudioDataset::writeCache()
{
    const auto file = getDatasetFile();
 
    DynamicObject* dataset = new DynamicObject();
    Array<var> rows;
 
    for (const auto& category : m_dataset)
    {
        const auto key = category.first;
        const auto results = category.second;
 
        for (const auto& element : results)
        {
            DynamicObject* metadata = new DynamicObject();
            metadata->setProperty("CatID", var(key));
 
            const auto path = element.first;
            const auto embedding = element.second;
            metadata->setProperty("Path", var(path));
 
            Array<var> features;
            for (auto value : embedding)
                features.add(value);
            metadata->setProperty("Embedding", features);
 
            rows.add(metadata);
            metadata = nullptr;
        }
    }
    dataset->setProperty("Version", var(m_version));
    dataset->setProperty("Dataset", rows);
 
    FileOutputStream stream(file);
    if (stream.openedOk())
    {
        // overwrite an existing file
        stream.setPosition(0);
        stream.truncate();
        JSON::writeToStream(stream, dataset);
        return true;
    }
 
    return false;
}
 
bool AudioDataset::readCache()
{
    const auto file = getDatasetFile();
    assert(file.existsAsFile());
 
    m_dataset.clear();
 
    auto json = JSON::parse(file);
    auto version = json.getProperty(Identifier("Version"), 0).toString().toStdString();
    if (version != m_version)
    {
        // treat read cache as fail to trigger running updated analysis
        return 0;
    }
 
    var result = json.getProperty(Identifier("Dataset"), 0);
    for (int i = 0; i < result.size(); ++i)
    {
        // Get the data from the cache file
        auto CatID = result[i].getProperty(Identifier("CatID"), 0).toString().toStdString();
        auto Path = result[i].getProperty(Identifier("Path"), 0).toString().toStdString();
        var array = result[i].getProperty(Identifier("Embedding"), 0);
        std::vector<float> embedding;
        for (int j = 0; j < array.size(); ++j)
            embedding.push_back(array[j]);
 
        auto metadata = std::make_pair(Path, embedding);
        m_dataset[CatID].push_back(metadata);
    }
 
    return 1;
}
 
bool AudioDataset::cacheExists() { return getDatasetFile().existsAsFile(); }
 
bool AudioDataset::deleteCache() { return getDatasetFile().deleteFile(); }
 
std::string AudioDataset::extractCatID(std::string filename) const
{
    std::string delimiter = "_";
    std::string key = filename.substr(0, filename.find(delimiter));
    key = fixKeyErrors(key);
    return key;
}
 
std::vector<std::pair<std::string, float>> AudioDataset::sample(std::string path, std::size_t k, bool descending) const
{
    assert(k > 0);
 
    std::vector<std::pair<std::string, float>> query_results;
 
    auto filename = juce::File(path).getFileName().toStdString();
    const auto key = extractCatID(filename);
    const auto success = std::binary_search(CatIDSet.begin(), CatIDSet.end(), key);
    if (success)
    {
        if (m_dataset.count(key) > 0)
        {
            //-------------------------
            //-------------------------
 
            const auto results = m_dataset.at(key);
            if (results.size() == 1)
            {
                // the query path is the only result
                query_results.push_back(std::make_pair(results.at(0).first, 1.f));
                return query_results;
            }
 
            //-------------------------
            //-------------------------
 
            const auto it = std::find_if(results.begin(), results.end(),
                                         [path](const std::pair<std::string, std::vector<float>>& element) {
                                             return element.first == path;
                                         });
            if (it == results.end())
            {
                // valid UCS CatID but path not found in dataset
                // file dragged in from outside factory assets?
                query_results.push_back(std::make_pair(path, 1.f));
                return query_results;
            }
            const auto query_index = std::distance(results.begin(), it);
            const auto query_features = results.at(query_index).second;
 
            //-------------------------
            //-------------------------
 
            const int scale = descending ? 1 : -1;
            std::vector<float> scores(results.size(), 0.f);
            for (std::size_t i = 0; i < results.size(); ++i)
            {
                const auto features = results.at(i).second;
                const auto score = std::inner_product(features.begin(), features.end(), query_features.begin(), 0.f);
                scores.at(i) = scale * score;
            }
 
            //-------------------------
            //-------------------------
 
            // find top-k matches based on cosine similarity score
            std::priority_queue<std::pair<float, std::size_t>, std::vector<std::pair<float, std::size_t>>,
                                std::greater<std::pair<float, std::size_t>>>
                    q;
 
            for (std::size_t i = 0; i < scores.size(); ++i)
            {
                if (q.size() < k) // if ((q.size() < k) && (scores[i] >= threshold))
                {
                    q.push(std::make_pair(scores.at(i), i));
                }
                else if (q.top().first < scores.at(i))
                {
                    q.pop();
                    q.push(std::make_pair(scores.at(i), i));
                }
            }
            assert(q.size() >= 1);
 
            //-------------------------
            //-------------------------
 
            while (!q.empty())
            {
                // smallest element at top
                const auto result = q.top();
                const auto score = scale * result.first;
                query_results.push_back(std::make_pair(results.at(result.second).first, score));
                q.pop();
            }
 
            return query_results;
        }
    }
 
    // invalid CatID or not in dataset
    query_results.push_back(std::make_pair(path, 1.f));
    return query_results;
}
 
File AudioDataset::getFactoryAssetDirectory() const
{
#ifdef JUCE_MAC
    return File::getSpecialLocation(File::commonApplicationDataDirectory).getFullPathName() +
           "/Application Support/Krotos/" + JucePlugin_Name + "/Factory Assets/";
#else
    return File("C:\\ProgramData\\Krotos\\" JucePlugin_Name "\\Factory Assets\\");
#endif
}
 
File AudioDataset::getDatasetFile() const
{
#ifdef JUCE_MAC
    return File::getSpecialLocation(File::commonApplicationDataDirectory).getFullPathName() +
           "/Application Support/Krotos/" + JucePlugin_Name + "/ttpResources/embeddings.json";
#else
    return File("C:\\ProgramData\\Krotos\\" JucePlugin_Name "\\ttpResources\\embeddings.json");
#endif
}
 
} // namespace krotos