video-pc/AvRecorder/encoder/audio_mixer.cpp

#include "audio_mixer.h"

#include "basic/basic.h"
#include "basic/frame.h"
#include "libavutil/error.h"

#include <QDebug>

#include <windows.h>

AVSampleFormat BitsToFmt(int bits)
{
    switch (bits) {
    case 8:
        return AV_SAMPLE_FMT_U8;
    case 16:
        return AV_SAMPLE_FMT_S16;
    case 32:
        return AV_SAMPLE_FMT_S32;
    case 64:
        return AV_SAMPLE_FMT_S64;
    default:
        return AV_SAMPLE_FMT_FLT;
    }
}

int FmtToBits(AVSampleFormat fmt)
{
    switch (fmt) {
    case AV_SAMPLE_FMT_U8:
        return 8;
    case AV_SAMPLE_FMT_S16:
        return 16;
    case AV_SAMPLE_FMT_S32:
        return 32;
    case AV_SAMPLE_FMT_S64:
        return 64;
    default:
        return 32;
    }
}

int SizeToNbSamples(int size, int bitsPerSample, int nbChannels)
{
    return (size << 3) / bitsPerSample / nbChannels;
}

int NbSamplesToSize(int nbSamples, int bitsPerSample, int nbChannels)
{
    return (nbSamples * bitsPerSample * nbChannels) >> 3;
}

bool FrameQueue::Init(int channelNums, int sampleRate, AVSampleFormat fmt, int nbSamples)
{
    _front = 0;
    _sampleRate = sampleRate;
    _fmt = fmt;
    _nbSamples = nbSamples;
    _usedLinesize = nbSamples * channelNums * (fmt == AV_SAMPLE_FMT_S16 ? 2 : 4);
    av_channel_layout_default(&_layout, channelNums);
    _queue.emplace(fmt, &_layout, sampleRate, nbSamples);
    return true;
}

Frame<MediaType::AUDIO> FrameQueue::Pop()
{
    if (_queue.size() > 1) {
        auto frame = std::move(_queue.front());
        _queue.pop();
        return frame;
    }
    return Frame<MediaType::AUDIO>();
}

void FrameQueue::Push(uint8_t* data, int length)
{
    if (length > _usedLinesize) { // 递归调用
        Push(data, length / 2);
        Push(data + length / 2, length / 2 + length % 2);
        return;
    }
    auto&& frame = _queue.back().frame;
    int secondLength = _front + length - _usedLinesize;
    if (secondLength <= 0) { // 第一段缓存是够用的
        memcpy(frame->data[0] + _front, data, length);
        _front += length;
        return;
    }
    // 第一段缓存不够用
    int firstLength = length - secondLength;
    if (firstLength > 0) {
        memcpy(frame->data[0] + _front, data, firstLength);
    }
    // 载入一段新缓存
    _queue.emplace(_fmt, &_layout, _sampleRate, _nbSamples);
    memcpy(_queue.back().frame->data[0], data + firstLength, secondLength);
    _front = secondLength;
}

bool Resampler::Open(int inChannelNums,
                     int inSampleRate,
                     AVSampleFormat inFmt,
                     int outChannelNums,
                     int outSampleRate,
                     AVSampleFormat outFmt,
                     int outNbSample)
{
    Close();
    _swrCtx = swr_alloc();
    if (!_swrCtx) {
        __DebugPrint("swr_alloc failed");
        return false;
    }

    AVChannelLayout tmpLayout;
    av_channel_layout_default(&tmpLayout, inChannelNums);
    av_opt_set_chlayout(_swrCtx, "in_chlayout", &tmpLayout, 0);
    av_opt_set_int(_swrCtx, "in_sample_rate", inSampleRate, 0);
    av_opt_set_sample_fmt(_swrCtx, "in_sample_fmt", inFmt, 0);
    if (!_fromQueue.Init(inChannelNums, inSampleRate, inFmt, inSampleRate / 100 * 2)) {
        __DebugPrint("_fromQueue.Init failed");
        return false;
    }

    av_channel_layout_default(&tmpLayout, outChannelNums);
    av_opt_set_chlayout(_swrCtx, "out_chlayout", &tmpLayout, 0);
    av_opt_set_int(_swrCtx, "out_sample_rate", outSampleRate, 0);
    av_opt_set_sample_fmt(_swrCtx, "out_sample_fmt", outFmt, 0);
    if (swr_init(_swrCtx) < 0) {
        Close();
        __DebugPrint("swr_init(_swrCtx) failed\n");
        return false;
    }
    if (!_toQueue.Init(outChannelNums, outSampleRate, outFmt, outNbSample)) {
        __DebugPrint("_toQueue.Init failed");
        Close();
        return false;
    }
    _swrFrame = Frame<MediaType::AUDIO>::Alloc(outFmt,
                                                           &tmpLayout,
                                                           outSampleRate,
                                                           outSampleRate / 100 * 2);
    if (!_swrFrame) {
        __DebugPrint("Alloc _swrFrame failed");
        Close();
        return false;
    }

    return true;
}

void Resampler::Close()
{
    Free(_swrCtx, [this] { swr_free(&_swrCtx); });
    Free(_swrFrame, [this] { av_frame_free(&_swrFrame); });
}

AVFrame* AudioMixer::Convert(uint32_t index, uint8_t* inBuf, uint32_t size)
{
    std::lock_guard<std::mutex> locker(_mutex);
    auto iter = _audioInputInfos.find(index);
    if (iter == _audioInputInfos.end()) {
        __DebugPrint("Audio input index %u not found", index);
        return nullptr;
    }

    // // 添加调试信息
    // static int debugCounter = 0;
    // if (++debugCounter % 100 == 0) {
    //     qDebug() << "AudioMixer::Convert - Input size:" << size << "bytes, Input index: " << index;
    // }

    if (!iter->second.resampler->Convert(inBuf, size)) {
        __DebugPrint("Resampler::Convert failed");
        return nullptr;
    }
    return _AdjustVolume() ? _outputFrame : nullptr;
}

bool AudioMixer::_AdjustVolume()
{
    // 检测所有流之间是不是相差太大了以及缓存的数据是不是太多了
    // 如果缓存的数据太多，直接将所有的队列删除同样的数据
    // 如果两个流直接数据相差太大，将多的那个减到和少的那个一样
    constexpr int MAX_DIFF = 10;
    constexpr int MAX_BUF_SIZE = 20;
    int minSize = INT_MAX;
    int maxSize = INT_MIN;
    FrameQueue* maxQueue = nullptr;
#undef min
    for (auto&& iter : _audioInputInfos) {
        auto&& queue = iter.second.resampler->GetQueue();
        if (queue.IsEmpty()) {
            return false;
        }
        minSize = std::min(minSize, (int) queue.GetSize());
        if (maxSize < (int) queue.GetSize()) {
            maxSize = (int) queue.GetSize();
            maxQueue = &queue;
        }
    }

    if (maxSize - minSize > MAX_DIFF) {
        // __DebugPrint("Clear MAX_DIFF");
        for (int i = 0; i < maxSize - minSize; ++i) {
            maxQueue->Pop();
        }
    }

    auto writeStream = (float*)(_outputFrame->data[0]);
    memset(writeStream, 0, _outputFrame->linesize[0]);

    for (auto iter = _audioInputInfos.begin(); iter != _audioInputInfos.end(); ++iter) {
        auto&& frameQueue = iter->second.resampler->GetQueue();
        if (minSize > MAX_BUF_SIZE) {
            // __DebugPrint("Clear MAX_BUF_SIZE");
            for (int i = 0; i < minSize - 2; ++i) {
                frameQueue.Pop();
            }
        }
        auto frame = frameQueue.Pop();
        auto scale = iter->second.scale;
        auto readStream = (float*) (frame.frame->data[0]);
        iter->second.volume = readStream[0] * scale;

        // 逐个计算赋值，优化音量处理
        for (int idx = 0; idx < _outputFrame->nb_samples; ++idx) {
            float sample = readStream[idx] * scale;
            // 使用软限制而不是硬限制，减少失真
            if (sample > 0.8f) {
                sample = 0.8f + (sample - 0.8f) * 0.3f; // 软限制
            } else if (sample < -0.8f) {
                sample = -0.8f + (sample + 0.8f) * 0.3f; // 软限制
            }
            writeStream[idx] += sample;
        }
    }
    return true;
}

AudioMixer::AudioMixer()
    : _inited(false)
{}

AudioMixer::~AudioMixer()
{
    // delete out_buf;
    if (_inited) {
        Close();
    }
}

bool AudioMixer::AddAudioInput(uint32_t index,
                               uint32_t sampleRate,
                               uint32_t channels,
                               uint32_t bitsPerSample,
                               AVSampleFormat format)
{
    std::lock_guard<std::mutex> locker(_mutex);
    if (_inited) {
        __DebugPrint("AddAudioInput called after init");
        return false;
    }
    // 根据index保存是否已经存在
    if (_audioInputInfos.find(index) != _audioInputInfos.end()) {
        __DebugPrint("AddAudioInput duplicate index %u", index);
        return false;
    }

    auto& filterInfo = _audioInputInfos[index];
    // 设置音频相关参数
    filterInfo.sampleRate = sampleRate;
    filterInfo.channels = channels;
    filterInfo.bitsPerSample = bitsPerSample;
    filterInfo.format = format;
    filterInfo.name = std::string("input") + std::to_string(index);
    return true;
}

bool AudioMixer::AddAudioOutput(const uint32_t sampleRate,
                                const uint32_t channels,
                                const uint32_t bitsPerSample,
                                const AVSampleFormat format)
{
    std::lock_guard<std::mutex> locker(_mutex);
    if (_inited) {
        __DebugPrint("AddAudioOutput called after init");
        return false;
    }
    // 设置音频相关参数
    _audioOutputInfo.sampleRate = sampleRate;
    _audioOutputInfo.channels = channels;
    _audioOutputInfo.bitsPerSample = bitsPerSample;
    _audioOutputInfo.format = format;
    _audioOutputInfo.name = "output";
    return true;
}

bool AudioMixer::SetOutFrameSize(int outFrameSize)
{
    if (_outFrameSize == outFrameSize) {
        return true;
    }
    _outFrameSize = outFrameSize;
    for (auto&& filterInfoPair : _audioInputInfos) {
        auto&& filterInfo = filterInfoPair.second;
        filterInfo.resampler = std::make_unique<Resampler>();
        if (!filterInfo.resampler->Open(filterInfo.channels,
                                               filterInfo.sampleRate,
                                               filterInfo.format,
                                               _audioOutputInfo.channels,
                                               _audioOutputInfo.sampleRate,
                                               _audioOutputInfo.format,
                                               outFrameSize)) {
            __DebugPrint("Resampler::Open failed for input %s", filterInfo.name.c_str());
            return false;
        }
    }
    AVChannelLayout tmpLayout;
    av_channel_layout_default(&tmpLayout, _audioOutputInfo.channels);
    Free(_outputFrame, [this] { av_frame_free(&_outputFrame); });
    _outputFrame = Frame<MediaType::AUDIO>::Alloc(_audioOutputInfo.format,
                                                              &tmpLayout,
                                                              _audioOutputInfo.sampleRate,
                                                              outFrameSize);
    if (!_outputFrame) {
        __DebugPrint("Alloc _outputFrame failed");
        return false;
    }
    _inited = true;
    return true;
}

bool AudioMixer::Close()
{
    if (!_inited) {
        return true;
    }
    _inited = false;
    std::lock_guard<std::mutex> locker(_mutex);
    _audioInputInfos.clear();
    Free(_outputFrame, [this] { av_frame_free(&_outputFrame); });
    _outFrameSize = 0;
    return true;
}

AudioMixer::AudioInfo* AudioMixer::GetInputInfo(uint32_t index)
{
    auto iter = _audioInputInfos.find(index);
    return iter == _audioInputInfos.end() ? nullptr : &(iter->second);
}

bool Resampler::Convert(uint8_t* data, int size)
{
    if (data == nullptr) {
        return false;
    }
    _fromQueue.Push(data, size);
    for (;;) { // 转换
        auto frame = _fromQueue.Pop();
        if (frame.frame == nullptr) {
            break;
        }
        int ret = swr_convert(_swrCtx,
                              _swrFrame->data,
                              _swrFrame->nb_samples,
                              (const uint8_t**) frame.frame->data,
                              frame.frame->nb_samples);
        if (ret < 0) {
            char buffer[AV_ERROR_MAX_STRING_SIZE];
            av_strerror(ret, buffer, sizeof(buffer));
            qDebug() << "swr_convert failed" << ret << QString::fromLatin1(buffer);
            return false;
        }

        _toQueue.Push(_swrFrame->data[0], _swrFrame->linesize[0]);
    }
    return true;
}