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

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