SpectralDenoiser.h

Go to the documentation of this file.

// DSPark - Professional Audio DSP Framework
// Copyright (c) 2026 Cristian Moresi - MIT License
 
#pragma once
 
#include "../Core/AudioBuffer.h"
#include "../Core/AudioSpec.h"
#include "../Core/DspMath.h"
#include "../Core/SpectralProcessor.h"
#include "../Core/StateBlob.h"
 
#include <algorithm>
#include <atomic>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <vector>
 
namespace dspark {
 
template <FloatType T>
class SpectralDenoiser
{
public:
    // -- Lifecycle ---------------------------------------------------------------
 
    void prepare(const AudioSpec& spec, int fftSize = 2048)
    {
        if (!spec.isValid()) return;
        prepared_.store(false, std::memory_order_relaxed);
        numChannels_ = spec.numChannels;
        stft_.prepare(spec, fftSize, fftSize / 4);
        numBins_ = stft_.getNumBins();
 
        profile_.assign(static_cast<size_t>(numBins_), 0.0f);
        gains_.assign(static_cast<size_t>(numChannels_),
                      std::vector<float>(static_cast<size_t>(numBins_), 1.0f));
        smooth_.assign(static_cast<size_t>(numBins_), 1.0f);
 
        prepared_.store(true, std::memory_order_relaxed);
        reset();
    }
 
    void reset() noexcept
    {
        if (!prepared_.load(std::memory_order_relaxed)) return;
        stft_.reset();
        for (auto& g : gains_)
            std::fill(g.begin(), g.end(), 1.0f);
        callCounter_ = 0;
    }
 
    void clearProfile() noexcept
    {
        std::fill(profile_.begin(), profile_.end(), 0.0f);
    }
 
    // -- Parameters (thread-safe) ---------------------------------------------------
 
    void setLearning(bool learning) noexcept
    {
        learning_.store(learning, std::memory_order_relaxed);
    }
 
    void setReduction(T db) noexcept
    {
        if (!std::isfinite(db)) return;
        reduction_.store(std::clamp(db, T(0), T(40)), std::memory_order_relaxed);
    }
 
    void setThreshold(T factor) noexcept
    {
        if (!std::isfinite(factor)) return;
        threshold_.store(std::clamp(factor, T(1), T(8)), std::memory_order_relaxed);
    }
 
    [[nodiscard]] bool getLearning() const noexcept { return learning_.load(std::memory_order_relaxed); }
    [[nodiscard]] T getReduction() const noexcept { return reduction_.load(std::memory_order_relaxed); }
    [[nodiscard]] T getThreshold() const noexcept { return threshold_.load(std::memory_order_relaxed); }
 
    [[nodiscard]] int getLatency() const noexcept { return stft_.getLatency(); }
 
    [[nodiscard]] std::vector<uint8_t> getState() const
    {
        StateWriter w(stateId("DNSE"), 1);
        // Explicit float casts: the blob stores float, and with T = double the
        // unqualified write(key, double) would be ambiguous (float/int32/bool).
        w.write("reduction", static_cast<float>(reduction_.load(std::memory_order_relaxed)));
        w.write("threshold", static_cast<float>(threshold_.load(std::memory_order_relaxed)));
        return w.blob();
    }
 
    bool setState(const uint8_t* data, size_t size)
    {
        StateReader r(data, size);
        if (!r.isValid() || r.processorId() != stateId("DNSE")) return false;
        setReduction(static_cast<T>(r.read("reduction", 18.0f)));
        setThreshold(static_cast<T>(r.read("threshold", 2.0f)));
        return true;
    }
 
    // -- Processing -------------------------------------------------------------------
 
    void processBlock(AudioBufferView<T> buffer) noexcept
    {
        if (!prepared_.load(std::memory_order_relaxed)) return;
 
        const bool learning = learning_.load(std::memory_order_relaxed);
        const float floorGain = std::pow(
            10.0f, -static_cast<float>(reduction_.load(std::memory_order_relaxed)) / 20.0f);
        const float thresh = static_cast<float>(threshold_.load(std::memory_order_relaxed));
 
        // The STFT invokes the callback once per PROCESSED channel per hop
        // (in channel order) - and it processes min(buffer, prepared)
        // channels. Modulo by that same effective count, reset per block, or
        // a narrow buffer over a wider spec would rotate the per-channel
        // gain memories between hops (channel 0 alternating onto channel 1's
        // release state - measured 0.004 divergence versus a mono-prepared
        // twin before this fix).
        const int nChEff = std::max(1, std::min(buffer.getNumChannels(), numChannels_));
        callCounter_ = 0;
 
        stft_.processBlock(buffer, [this, learning, floorGain, thresh, nChEff](T* bins, int numBins)
        {
            auto& chGain = gains_[static_cast<size_t>(callCounter_ % nChEff)];
            ++callCounter_;
 
            // Per-bin raw gate decision.
            for (int k = 0; k < numBins; ++k)
            {
                const float re = static_cast<float>(bins[2 * k]);
                const float im = static_cast<float>(bins[2 * k + 1]);
                const float mag = std::sqrt(re * re + im * im);
 
                if (learning)
                {
                    // Profile: peak-hold with a gentle average pull-up.
                    auto& p = profile_[static_cast<size_t>(k)];
                    p = std::max(p * 0.995f + mag * 0.005f, std::max(p, mag * 0.8f));
                }
 
                const float open = profile_[static_cast<size_t>(k)] * thresh;
                smooth_[static_cast<size_t>(k)] = (mag > open) ? 1.0f : floorGain;
            }
 
            // Frequency smoothing (3-bin average) defends against isolated
            // flickering bins - the source of musical noise.
            for (int k = 0; k < numBins; ++k)
            {
                const float a = smooth_[static_cast<size_t>(std::max(k - 1, 0))];
                const float b = smooth_[static_cast<size_t>(k)];
                const float c = smooth_[static_cast<size_t>(std::min(k + 1, numBins - 1))];
                float target = (a + b + c) * (1.0f / 3.0f);
 
                // Temporal smoothing: instant attack (open fast), slow release.
                auto& g = chGain[static_cast<size_t>(k)];
                g = (target > g) ? target : (g * 0.85f + target * 0.15f);
 
                bins[2 * k] = static_cast<T>(static_cast<float>(bins[2 * k]) * g);
                bins[2 * k + 1] = static_cast<T>(static_cast<float>(bins[2 * k + 1]) * g);
            }
        });
    }
 
private:
    SpectralProcessor<T> stft_;
    int numChannels_ = 0;
    int numBins_ = 0;
    std::atomic<bool> prepared_ { false };
 
    std::vector<float> profile_;               
    std::vector<std::vector<float>> gains_;    
    std::vector<float> smooth_;                
    int callCounter_ = 0;
 
    std::atomic<bool> learning_ { false };
    std::atomic<T> reduction_ { T(18) };
    std::atomic<T> threshold_ { T(2) };
};
 
} // namespace dspark