git.sesse.net Git - nageru/blob - nageru/audio_encoder.cpp

   1 #include "audio_encoder.h"
   2
   3 extern "C" {
   4 #include <libavcodec/avcodec.h>
   5 #include <libavformat/avformat.h>
   6 #include <libswresample/swresample.h>
   7 #include <libavutil/channel_layout.h>
   8 #include <libavutil/error.h>
   9 #include <libavutil/frame.h>
  10 #include <libavutil/mem.h>
  11 #include <libavutil/opt.h>
  12 #include <libavutil/rational.h>
  13 #include <libavutil/samplefmt.h>
  14 }
  15
  16 #include <assert.h>
  17 #include <errno.h>
  18 #include <stdio.h>
  19 #include <stdlib.h>
  20 #include <memory>
  21 #include <string>
  22 #include <vector>
  23
  24 #include "defs.h"
  25 #include "shared/mux.h"
  26 #include "shared/shared_defs.h"
  27 #include "shared/timebase.h"
  28
  29 using namespace std;
  30
  31 AudioEncoder::AudioEncoder(const string &codec_name, int bit_rate, const AVOutputFormat *oformat)
  32 {
  33         const AVCodec *codec = avcodec_find_encoder_by_name(codec_name.c_str());
  34         if (codec == nullptr) {
  35                 fprintf(stderr, "ERROR: Could not find codec '%s'\n", codec_name.c_str());
  36                 abort();
  37         }
  38
  39         ctx = avcodec_alloc_context3(codec);
  40         ctx->bit_rate = bit_rate;
  41         ctx->sample_rate = OUTPUT_FREQUENCY;
  42         ctx->sample_fmt = codec->sample_fmts[0];
  43         ctx->ch_layout.order = AV_CHANNEL_ORDER_NATIVE;
  44         ctx->ch_layout.nb_channels = 2;
  45         ctx->ch_layout.u.mask = AV_CH_LAYOUT_STEREO;
  46         ctx->time_base = AVRational{1, TIMEBASE};
  47         if (oformat->flags & AVFMT_GLOBALHEADER) {
  48                 ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
  49         }
  50         if (avcodec_open2(ctx, codec, NULL) < 0) {
  51                 fprintf(stderr, "Could not open codec '%s'\n", codec_name.c_str());
  52                 abort();
  53         }
  54
  55         resampler = nullptr;
  56         int ok = swr_alloc_set_opts2(&resampler,
  57                                      /*out_ch_layout=*/&ctx->ch_layout,
  58                                      /*out_sample_fmt=*/ctx->sample_fmt,
  59                                      /*out_sample_rate=*/OUTPUT_FREQUENCY,
  60                                      /*in_ch_layout=*/&ctx->ch_layout,
  61                                      /*in_sample_fmt=*/AV_SAMPLE_FMT_FLT,
  62                                      /*in_sample_rate=*/OUTPUT_FREQUENCY,
  63                                      /*log_offset=*/0,
  64                                      /*log_ctx=*/nullptr);
  65         if (ok != 0) {
  66                 fprintf(stderr, "Allocating resampler failed.\n");
  67                 abort();
  68         }
  69
  70         if (swr_init(resampler) < 0) {
  71                 fprintf(stderr, "Could not open resample context.\n");
  72                 abort();
  73         }
  74
  75         audio_frame = av_frame_alloc();
  76 }
  77
  78 AudioEncoder::~AudioEncoder()
  79 {
  80         av_frame_free(&audio_frame);
  81         swr_free(&resampler);
  82         avcodec_free_context(&ctx);
  83 }
  84
  85 void AudioEncoder::encode_audio(const vector<float> &audio, int64_t audio_pts)
  86 {
  87         if (ctx->frame_size == 0) {
  88                 // No queueing needed.
  89                 assert(audio_queue.empty());
  90                 assert(audio.size() % 2 == 0);
  91                 encode_audio_one_frame(&audio[0], audio.size() / 2, audio_pts);
  92                 return;
  93         }
  94
  95         int64_t sample_offset = audio_queue.size();
  96
  97         audio_queue.insert(audio_queue.end(), audio.begin(), audio.end());
  98         size_t sample_num;
  99         for (sample_num = 0;
 100              sample_num + ctx->frame_size * 2 <= audio_queue.size();
 101              sample_num += ctx->frame_size * 2) {
 102                 int64_t adjusted_audio_pts = audio_pts + (int64_t(sample_num) - sample_offset) * TIMEBASE / (OUTPUT_FREQUENCY * 2);
 103                 encode_audio_one_frame(&audio_queue[sample_num],
 104                                        ctx->frame_size,
 105                                        adjusted_audio_pts);
 106         }
 107         audio_queue.erase(audio_queue.begin(), audio_queue.begin() + sample_num);
 108
 109         last_pts = audio_pts + audio.size() * TIMEBASE / (OUTPUT_FREQUENCY * 2);
 110 }
 111
 112 void AudioEncoder::encode_audio_one_frame(const float *audio, size_t num_samples, int64_t audio_pts)
 113 {
 114         audio_frame->pts = audio_pts;
 115         audio_frame->nb_samples = num_samples;
 116         audio_frame->ch_layout.order = AV_CHANNEL_ORDER_NATIVE;
 117         audio_frame->ch_layout.nb_channels = 2;
 118         audio_frame->ch_layout.u.mask = AV_CH_LAYOUT_STEREO;
 119         audio_frame->format = ctx->sample_fmt;
 120         audio_frame->sample_rate = OUTPUT_FREQUENCY;
 121
 122         if (av_samples_alloc(audio_frame->data, nullptr, 2, num_samples, ctx->sample_fmt, 0) < 0) {
 123                 fprintf(stderr, "Could not allocate %zu samples.\n", num_samples);
 124                 abort();
 125         }
 126
 127         if (swr_convert(resampler, audio_frame->data, num_samples, reinterpret_cast<const uint8_t **>(&audio), num_samples) < 0) {
 128                 fprintf(stderr, "Audio conversion failed.\n");
 129                 abort();
 130         }
 131
 132         int err = avcodec_send_frame(ctx, audio_frame);
 133         if (err < 0) {
 134                 fprintf(stderr, "avcodec_send_frame() failed with error %d\n", err);
 135                 abort();
 136         }
 137
 138         for ( ;; ) {  // Termination condition within loop.
 139                 AVPacketWithDeleter pkt = av_packet_alloc_unique();
 140                 pkt->data = nullptr;
 141                 pkt->size = 0;
 142                 int err = avcodec_receive_packet(ctx, pkt.get());
 143                 if (err == 0) {
 144                         pkt->stream_index = 1;
 145                         pkt->flags = 0;
 146                         for (Mux *mux : muxes) {
 147                                 mux->add_packet(*pkt, pkt->pts, pkt->dts);
 148                         }
 149                 } else if (err == AVERROR(EAGAIN)) {
 150                         break;
 151                 } else {
 152                         fprintf(stderr, "avcodec_receive_frame() failed with error %d\n", err);
 153                         abort();
 154                 }
 155         }
 156
 157         av_freep(&audio_frame->data[0]);
 158         av_frame_unref(audio_frame);
 159 }
 160
 161 void AudioEncoder::encode_last_audio()
 162 {
 163         if (!audio_queue.empty()) {
 164                 // Last frame can be whatever size we want.
 165                 assert(audio_queue.size() % 2 == 0);
 166                 encode_audio_one_frame(&audio_queue[0], audio_queue.size() / 2, last_pts);
 167                 audio_queue.clear();
 168         }
 169
 170         if (ctx->codec->capabilities & AV_CODEC_CAP_DELAY) {
 171                 // Collect any delayed frames.
 172                 for ( ;; ) {
 173                         AVPacketWithDeleter pkt = av_packet_alloc_unique();
 174                         pkt->data = nullptr;
 175                         pkt->size = 0;
 176                         int err = avcodec_receive_packet(ctx, pkt.get());
 177                         if (err == 0) {
 178                                 pkt->stream_index = 1;
 179                                 pkt->flags = 0;
 180                                 for (Mux *mux : muxes) {
 181                                         mux->add_packet(*pkt, pkt->pts, pkt->dts);
 182                                 }
 183                         } else if (err == AVERROR_EOF) {
 184                                 break;
 185                         } else {
 186                                 fprintf(stderr, "avcodec_receive_frame() failed with error %d\n", err);
 187                                 abort();
 188                         }
 189                 }
 190         }
 191 }
 192
 193 AVCodecParametersWithDeleter AudioEncoder::get_codec_parameters()
 194 {
 195         AVCodecParameters *codecpar = avcodec_parameters_alloc();
 196         avcodec_parameters_from_context(codecpar, ctx);
 197         return AVCodecParametersWithDeleter(codecpar);
 198 }