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

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