libav transcoder prototype

#include <stdio.h>
#include <stdlib.h>

#include <libavformat/avformat.h>
#include <libavutil/channel_layout.h>
#include <libavutil/error.h>
#include <libavutil/avstring.h>

#include <libswscale/swscale.h>

#define DUMP_FRAMES 0

#define BUFFER_SIZE 1024

#define AUDIO_OUT_CODEC AV_CODEC_ID_AAC
#define AUDIO_OUT_CHANNELS 2
#define AUDIO_OUT_BITRATE 96000

#define VIDEO_OUT_CODEC AV_CODEC_ID_H264
#define VIDEO_OUT_BITRATE 400000

#define VIDEO_CONTAINER_SIZE 16777216  // 16 MB

typedef struct {
  AVStream * stream;
  AVCodecContext * in_codec_ctx;
  AVCodecContext * out_codec_ctx;
} TranscodingContext;

static char *const get_error_text(const int error) {
  static char error_buffer[255];

  av_strerror(error, error_buffer, sizeof(error_buffer));

  return error_buffer;
}

void ok(void) {
  fprintf(stderr, "OK\n");
}

int failed(int error) {
  fprintf(stderr, "FAILED: %s\n", get_error_text(error));
  return error;
}

int open_input_file(const char *filename, AVFormatContext **in_fmt_ctx) {
  int error = 0;

  fprintf(stderr, "Opening input file %s ... ", filename);
  if ((error = avformat_open_input(in_fmt_ctx, filename, NULL, NULL)) >= 0) {
    ok();
  } else {
    failed(error);
    return error;
  }

  fprintf(stderr, "Finding stream info... ");
  if ((error = avformat_find_stream_info(*in_fmt_ctx, NULL)) >= 0) {
    ok();
  } else {
    return failed(error);
  }

  return 0;
}

void print_codec(TranscodingContext * trans_ctx, AVCodec * codec) {
  AVCodecContext * stream_codec = trans_ctx->stream->codec;
  enum AVMediaType codec_type = stream_codec->codec_type;

  const AVCodec * c = codec ? codec :
    trans_ctx->in_codec_ctx ? trans_ctx->in_codec_ctx->codec
    : NULL;

  fprintf(stderr, "%s (%s)",
          (codec_type == AVMEDIA_TYPE_VIDEO ? "video" :
           codec_type == AVMEDIA_TYPE_AUDIO ? "audio" :
           codec_type == AVMEDIA_TYPE_SUBTITLE ? "subtitle" :
           codec_type == AVMEDIA_TYPE_DATA ? "data" :
           "unknown"),
          (c ? c->name : "no codec"));
}

AVCodecContext * copy_codec_context(AVCodec * codec, AVCodecContext * src_ctx) {
  int error;

  AVCodecContext * new_ctx = avcodec_alloc_context3(codec);
  fprintf(stderr, "      Copying codec context... ");
  if ((error = avcodec_copy_context(new_ctx, src_ctx)) >= 0) {
    ok();
  } else {
    failed(error);
    return NULL;
  }

  return new_ctx;
}

TranscodingContext * open_codec(AVStream * stream) {
  TranscodingContext * trans_ctx = calloc(1, sizeof(TranscodingContext));
  trans_ctx->stream = stream;

  int error = 0;

  fprintf(stderr, " [%02d] ", stream->index);

  fprintf(stderr, "Finding decoder... ");
  AVCodec * codec = avcodec_find_decoder(stream->codec->codec_id);
  print_codec(trans_ctx, codec);
  fprintf(stderr, " ");

  if (codec) {
    ok();
  } else {
    fprintf(stderr, "FAILED\n");
    return trans_ctx;
  }

  AVCodecContext * in_codec_ctx = copy_codec_context(codec, stream->codec);
  if (!in_codec_ctx)
    return trans_ctx;

  fprintf(stderr, "      Opening codec... ");
  in_codec_ctx->refcounted_frames = 1;
  if ((error = avcodec_open2(in_codec_ctx, codec, NULL)) >= 0) {
    trans_ctx->in_codec_ctx = in_codec_ctx;
    ok();
  } else {
    failed(error);
    return trans_ctx;
  }

  return trans_ctx;
}

void decode(AVFormatContext * in_fmt_ctx, TranscodingContext ** stream_contexts, AVFrame *** frames, int * num_frames) {
  AVPacket packet;
  AVFrame * frame = av_frame_alloc();

  int error = 0;
  int packet_num = 0;

  fprintf(stderr, "Reading packets\n");

  while (1) {
    av_init_packet(&packet);

    fprintf(stderr, " Packet %03d ", packet_num);

    if ((error = av_read_frame(in_fmt_ctx, &packet)) < 0) {
      if (error == AVERROR_EOF) {
        fprintf(stderr, "end of file\n");

        break;
      } else {
        fprintf(stderr, "FAILED: %s", get_error_text(error));

        continue;
      }
    }

    int len, got_frame;
    TranscodingContext * trans_ctx = stream_contexts[packet.stream_index];
    AVCodecContext * codec = trans_ctx->in_codec_ctx;
    int offset = num_frames[packet.stream_index];
    int codec_type = codec ? codec->codec_type : 0;

    fprintf(stderr, "stream %02d ", packet.stream_index);
    print_codec(trans_ctx, NULL);
    fprintf(stderr, " ");

    if (!codec) {
      fprintf(stderr, "SKIPPING PACKET: no codec\n");
      continue;
    }

    // TODO
    if (codec_type != AVMEDIA_TYPE_VIDEO && codec_type != AVMEDIA_TYPE_AUDIO) {
      av_frame_free(&frame);
      continue;
    }

    while (packet.size > 0) {
      if (codec_type == AVMEDIA_TYPE_VIDEO) {
        len = avcodec_decode_video2(codec, frame, &got_frame, &packet);
      } else if (codec_type == AVMEDIA_TYPE_AUDIO) {
        len = avcodec_decode_audio4(codec, frame, &got_frame, &packet);
      }

      if (len < 0) {
        failed(len);
        break;
      }

      if (got_frame) {
        frames[packet.stream_index][offset] = frame;
        num_frames[packet.stream_index]++;
      } else {
        av_frame_free(&frame);
      }

      packet.size -= len;
      packet.data += len;

      frame = av_frame_alloc();

      ok();
    }

    packet_num++;
  }

  av_free_packet(&packet);
}

int check_sample_fmt(AVCodec *codec, enum AVSampleFormat sample_fmt) {
  const enum AVSampleFormat *p = codec->sample_fmts;

  while (*p != AV_SAMPLE_FMT_NONE) {
    if (*p == sample_fmt)
      return 1;
    p++;
  }

  return 0;
}

int buffer_offset(const AVIOContext * io_ctx) {
  return io_ctx->buf_ptr - io_ctx->buffer;
}

void encode(int stream_num, TranscodingContext * trans_ctx, AVFrame ** frames, int num_frames) {
  fprintf(stderr, " [%02d] ", stream_num);
  print_codec(trans_ctx, NULL);
  fprintf(stderr, "\n");

  AVCodecContext * in_codec_ctx = trans_ctx->in_codec_ctx;
  if (!in_codec_ctx) {
    fprintf(stderr, "      No codec; skipping\n");
    return;
  }

  enum AVMediaType codec_type = in_codec_ctx->codec_type;

  if (codec_type != AVMEDIA_TYPE_VIDEO && codec_type != AVMEDIA_TYPE_AUDIO) {
    fprintf(stderr, "      Not audio or video; skipping\n");
  } else {
    AVCodec * out_codec = NULL;
    AVFormatContext * out_fmt_ctx = avformat_alloc_context();
    AVStream * stream = NULL;

    uint16_t * samples = NULL;

    char filename[256];
    FILE * outfile = NULL;

    int error = 0;

    sprintf(filename, "/tmp/%s-%02d.mp4",
            (codec_type == AVMEDIA_TYPE_AUDIO ? "audio" : "video"), stream_num);

    fprintf(stderr, "      Guessing output format... ");
    if (!(out_fmt_ctx->oformat = av_guess_format(NULL, filename, NULL))) {
      fprintf(stderr, "FAILED\n");
      goto cleanup_encoder;
    }
    ok();

    av_strlcpy(out_fmt_ctx->filename, filename, sizeof(out_fmt_ctx->filename));

    fprintf(stderr, "      Finding output encoder... ");
    if (!(out_codec = avcodec_find_encoder(codec_type == AVMEDIA_TYPE_AUDIO ? AUDIO_OUT_CODEC : VIDEO_OUT_CODEC))) {
      fprintf(stderr, "FAILED\n");
      goto cleanup_encoder;
    }
    ok();

    stream = avformat_new_stream(out_fmt_ctx, out_codec);

    trans_ctx->out_codec_ctx = stream->codec;
    if (codec_type == AVMEDIA_TYPE_AUDIO) {
      stream->time_base.den = in_codec_ctx->sample_rate;
      stream->time_base.num = 1;

      trans_ctx->out_codec_ctx->channels = AUDIO_OUT_CHANNELS;
      trans_ctx->out_codec_ctx->channel_layout = av_get_default_channel_layout(AUDIO_OUT_CHANNELS);
      trans_ctx->out_codec_ctx->sample_rate = in_codec_ctx->sample_rate;
      trans_ctx->out_codec_ctx->sample_fmt = out_codec->sample_fmts[0];
      trans_ctx->out_codec_ctx->bit_rate = AUDIO_OUT_BITRATE;
      trans_ctx->out_codec_ctx->strict_std_compliance = FF_COMPLIANCE_EXPERIMENTAL;
      if (out_fmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
        trans_ctx->out_codec_ctx->flags |= CODEC_FLAG_GLOBAL_HEADER;

      fprintf(stderr, "      Checking sample format support... ");
      if (check_sample_fmt(out_codec, in_codec_ctx->sample_fmt)) {
        ok();
      } else {
        fprintf(stderr, "FAILED (resampling required)\n");
        goto cleanup_encoder;
        // TODO: handle resampling
      }

      /* frame->nb_samples = trans_ctx->out_codec_ctx->frame_size; */
      /* frame->format = trans_ctx->out_codec_ctx->sample_fmt; */
      /* frame->channel_layout = trans_ctx->out_codec_ctx->channel_layout; */
    } else {
      stream->time_base = in_codec_ctx->time_base;

      trans_ctx->out_codec_ctx->bit_rate = VIDEO_OUT_BITRATE;
      trans_ctx->out_codec_ctx->width = in_codec_ctx->width;
      trans_ctx->out_codec_ctx->height = in_codec_ctx->height;
      trans_ctx->out_codec_ctx->gop_size = 10;
      trans_ctx->out_codec_ctx->max_b_frames = 1;
      trans_ctx->out_codec_ctx->pix_fmt = AV_PIX_FMT_YUV420P;
      if (out_fmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
        trans_ctx->out_codec_ctx->flags |= CODEC_FLAG_GLOBAL_HEADER;
    }

    fprintf(stderr, "      Opening output codec... ");
    if ((error = avcodec_open2(trans_ctx->out_codec_ctx, out_codec, NULL)) < 0) {
      failed(error);
      goto cleanup_encoder;
    }
    ok();

    AVIOContext * io_ctx = NULL;
    unsigned char * io_buf = NULL;

    if (codec_type == AVMEDIA_TYPE_AUDIO) {
      // TODO
    } else {  // video
      AVPacket packet;
      int got_output;

      fprintf(stderr, "      Opening output container... ");
      io_buf = av_malloc(VIDEO_CONTAINER_SIZE);
      if (!io_buf) {
        fprintf(stderr, "FAILED: could not allocate buffer\n");
        goto cleanup_encoder;
      }        

      io_ctx = avio_alloc_context(io_buf, VIDEO_CONTAINER_SIZE, 1, NULL, NULL, NULL, NULL);
      if (!io_ctx) {
        fprintf(stderr, "FAILED: could not allocate IO context\n");
        goto cleanup_encoder;
      }

      io_ctx->seekable = AVIO_SEEKABLE_NORMAL;
      out_fmt_ctx->pb = io_ctx;
      ok();
      
      fprintf(stderr, "      Writing header starting at buffer offset %d... ", buffer_offset(io_ctx));
      if ((error = avformat_write_header(out_fmt_ctx, NULL)) < 0) {
        failed(error);
        goto cleanup_encoder;
      }
      ok();
      
      fprintf(stderr, "      Encoding frames\n");

      for (int i = 0; i < num_frames; i++) {
        fprintf(stderr, "        Frame %d, buffer offset %d... ", i, buffer_offset(io_ctx));

        av_init_packet(&packet);
        packet.size = 0;
        packet.data = NULL;

        if ((error = avcodec_encode_video2(trans_ctx->out_codec_ctx, &packet, frames[i], &got_output)) < 0) {
          fprintf(stderr, "FAILED to encode video: %s\n", get_error_text(error));
          continue;
        }

        if ((error = av_interleaved_write_frame(out_fmt_ctx, &packet)) < 0) {
          fprintf(stderr, "FAILED to write frame: %s\n", get_error_text(error));
          continue;
        }
        
        ok();
      }
      
      fprintf(stderr, "      Writing trailer starting at buffer offset %d... ", buffer_offset(io_ctx));
      if ((error = av_write_trailer(out_fmt_ctx)) < 0) {
        failed(error);
        goto cleanup_encoder;
      }
      ok();

      fprintf(stderr, "      Opening %s ...", filename);
      if (!(outfile = fopen(filename, "wb"))) {
        fprintf(stderr, "FAILED\n");
        goto cleanup_encoder;
      }
      ok();

      fwrite(io_ctx->buffer, 1, (io_ctx->buf_end - io_ctx->buffer), outfile);
    }


  cleanup_encoder:
    if (io_buf)
      av_free(io_buf);

    if (io_ctx)
      av_free(io_ctx);

    if (outfile)
      fclose(outfile);

    if (out_fmt_ctx)
      avformat_free_context(out_fmt_ctx);
  }
}

void dump_frames(int stream_num, TranscodingContext * trans_ctx, AVFrame ** frames, int num_frames) {
  AVCodecContext * in_codec_ctx = trans_ctx->in_codec_ctx;

  fprintf(stderr, " [%02d] ", stream_num);
  print_codec(trans_ctx, NULL);

  if (in_codec_ctx && in_codec_ctx->codec_type == AVMEDIA_TYPE_VIDEO) {
    fprintf(stderr, " dumping %d frames ", num_frames);
  } else {
    fprintf(stderr, " skipping\n");
    return;
  }

  AVFrame * rgb_frame = av_frame_alloc();

  int width = in_codec_ctx->width;
  int height = in_codec_ctx->height;

  fprintf(stderr, "(%dx%d)\n", width, height);

  int num_bytes = avpicture_get_size(PIX_FMT_RGB24, width, height);
  uint8_t * buffer = (uint8_t *)av_malloc(num_bytes * sizeof(uint8_t));

  avpicture_fill((AVPicture *)rgb_frame, buffer, PIX_FMT_RGB24, width, height);

  struct SwsContext * sws_ctx = sws_getContext(width, height, in_codec_ctx->pix_fmt,
                                               width, height, PIX_FMT_RGB24,
                                               SWS_BILINEAR, NULL, NULL, NULL);

  char filename[64];
  for (int i = 0; i < num_frames; i++) {
    sws_scale(sws_ctx, (uint8_t const * const *)frames[i]->data,
              frames[i]->linesize, 0, height, rgb_frame->data, rgb_frame->linesize);

    FILE * f = NULL;

    sprintf(filename, "/tmp/stream-%02d-frame-%03d.ppm", stream_num, i);

    fprintf(stderr, "      Frame %02d dumping to file %s ... ", i, filename);

    if (!(f = fopen(filename, "wb"))) {
      fprintf(stderr, "FAILED: cannot open file\n");
      continue;
    }

    fprintf(f, "P6\n%d %d\n255\n", width, height);

    for (int y = 0; y < height; y++)
      fwrite(rgb_frame->data[0] + y * rgb_frame->linesize[0], 1, width * 3, f);

    fclose(f);
    ok();
  }

  av_free(buffer);
  av_free(rgb_frame);
}

int main(int argc, char * argv[]) {
  if (argc < 2) {
    fprintf(stderr, "Usage: %s INPUT_FILE\n", argv[0]);
    return 1;
  }

  char * in_filename = argv[1];

  int error = 0;

  int num_streams = 0;
  TranscodingContext ** stream_contexts = NULL;
  AVFrame *** frames = NULL;
  int * num_frames = NULL;

  av_register_all();

  AVFormatContext * in_fmt_ctx = NULL;
  if ((error = open_input_file(in_filename, &in_fmt_ctx)) != 0)
    goto cleanup;

  num_streams = in_fmt_ctx->nb_streams;
  fprintf(stderr, "Found %d streams\n", num_streams);

  stream_contexts = calloc(num_streams, sizeof(TranscodingContext *));
  frames = calloc(num_streams, sizeof(char *));
  num_frames = calloc(num_streams, sizeof(int));

  for (int i = 0; i < num_streams; i++) {
    AVStream * stream = in_fmt_ctx->streams[i];
    stream_contexts[stream->index] = open_codec(stream);
    frames[stream->index] = calloc(BUFFER_SIZE, sizeof(AVFrame *));
    num_frames[stream->index] = 0;
  }

  decode(in_fmt_ctx, stream_contexts, frames, num_frames);

  fprintf(stderr, "Decoding successful\n");

  for (int i = 0; i < num_streams; i++) {
    fprintf(stderr, " [%02d] ", i);
    print_codec(stream_contexts[i], NULL);
    fprintf(stderr, " %d frames\n", num_frames[i]);
  }

  if (DUMP_FRAMES) {
    fprintf(stderr, "Dumping frames\n");
    for (int i = 0; i < num_streams; i++) {
      dump_frames(i, stream_contexts[i], frames[i], num_frames[i]);
    }
  }

  fprintf(stderr, "Encoding\n");
  for (int i = 0; i < num_streams; i++) {
    encode(i, stream_contexts[i], frames[i], num_frames[i]);
  }

 cleanup:
  for (int i = 0; i < num_streams; i++) {
    fprintf(stderr, " [%02d] Cleaning up\n", i);

    if (frames) {
      fprintf(stderr, "      Buffer... ");
      for (int j = 0; j < num_frames[i]; j++)
        av_frame_unref(frames[i][j]);
      ok();

      free(frames[i]);
    }

    if (stream_contexts && stream_contexts[i]) {
      if (stream_contexts[i]->in_codec_ctx) {
        fprintf(stderr, "      Input codec... ");
        avcodec_close(stream_contexts[i]->in_codec_ctx);
        ok();
      }
      if (stream_contexts[i]->out_codec_ctx) {
        fprintf(stderr, "      Output codec... ");
        avcodec_close(stream_contexts[i]->out_codec_ctx);
        ok();
      }

      free(stream_contexts[i]);
    }

    avcodec_close(in_fmt_ctx->streams[i]->codec);
  }

  if (stream_contexts)
    free(stream_contexts);
  if (frames)
    free(frames);
  if (num_frames)
    free(num_frames);

  exit(0); // TODO: fix double-free with valgrind

  if (in_fmt_ctx) {
    fprintf(stderr, "Closing input format context... ");
    avformat_close_input(&in_fmt_ctx);
    ok();
  }
}