Add encoder

Makefile

@@ -6,6 +6,7 @@
 XAV_DIR = c_src/xav
 PRIV_DIR = $(MIX_APP_PATH)/priv
 XAV_DECODER_SO = $(PRIV_DIR)/libxavdecoder.so
+XAV_ENCODER_SO = $(PRIV_DIR)/libxavencoder.so
 XAV_READER_SO = $(PRIV_DIR)/libxavreader.so
 XAV_VIDEO_CONVERTER_SO = $(PRIV_DIR)/libxavvideoconverter.so
 
@@ -15,6 +16,9 @@ XAV_VIDEO_CONVERTER_SO = $(PRIV_DIR)/libxavvideoconverter.so
 DECODER_HEADERS = $(XAV_DIR)/xav_decoder.h $(XAV_DIR)/decoder.h $(XAV_DIR)/video_converter.h $(XAV_DIR)/audio_converter.h $(XAV_DIR)/utils.h $(XAV_DIR)/channel_layout.h
 DECODER_SOURCES = $(XAV_DIR)/xav_decoder.c $(XAV_DIR)/decoder.c $(XAV_DIR)/video_converter.c $(XAV_DIR)/audio_converter.c $(XAV_DIR)/utils.c
 
+ENCODER_HEADERS = $(XAV_DIR)/xav_encoder.h $(XAV_DIR)/encoder.h $(XAV_DIR)/utils.h
+ENCODER_SOURCES = $(XAV_DIR)/xav_encoder.c $(XAV_DIR)/encoder.c $(XAV_DIR)/utils.c
+
 READER_HEADERS = $(XAV_DIR)/xav_reader.h $(XAV_DIR)/reader.h $(XAV_DIR)/video_converter.h $(XAV_DIR)/audio_converter.h $(XAV_DIR)/utils.h $(XAV_DIR)/channel_layout.h
 READER_SOURCES = $(XAV_DIR)/xav_reader.c $(XAV_DIR)/reader.c $(XAV_DIR)/video_converter.c $(XAV_DIR)/audio_converter.c $(XAV_DIR)/utils.c
 
@@ -42,7 +46,7 @@ ifneq (,$(wildcard /etc/fedora-release))
 	LFLAGS += $$(pkg-config --libs-only-L libavcodec libswscale libavutil libavformat libavdevice libswresample)
 endif
 
-all: $(XAV_DECODER_SO) $(XAV_READER_SO) $(XAV_VIDEO_CONVERTER_SO)
+all: $(XAV_DECODER_SO) $(XAV_READER_SO) $(XAV_VIDEO_CONVERTER_SO) $(XAV_ENCODER_SO)
 
 $(XAV_DECODER_SO): Makefile $(DECODER_SOURCES) $(DECODER_HEADERS)
 	mkdir -p $(PRIV_DIR)
@@ -56,6 +60,11 @@ $(XAV_VIDEO_CONVERTER_SO): Makefile $(VIDEO_CONVERTER_SOURCES) $(VIDEO_CONVERTER
 	mkdir -p $(PRIV_DIR)
 	$(CC) $(CFLAGS) $(IFLAGS) $(LFLAGS) $(VIDEO_CONVERTER_SOURCES) -o $(XAV_VIDEO_CONVERTER_SO) $(LDFLAGS)
 
+
+$(XAV_ENCODER_SO): Makefile $(ENCODER_SOURCES) $(ENCODER_HEADERS)
+	mkdir -p $(PRIV_DIR)
+	$(CC) $(CFLAGS) $(IFLAGS) $(LFLAGS) $(ENCODER_SOURCES) -o $(XAV_ENCODER_SO) $(LDFLAGS)
+
 format:
 	clang-format -i $(XAV_DIR)/*
 

c_src/xav/encoder.c

@@ -0,0 +1,107 @@
+#include "encoder.h"
+
+struct Encoder *encoder_alloc() {
+  struct Encoder *encoder = XAV_ALLOC(sizeof(struct Encoder));
+  encoder->c = NULL;
+  encoder->codec = NULL;
+  encoder->num_packets = 0;
+  encoder->max_num_packets = 8;
+  encoder->packets = XAV_ALLOC(encoder->max_num_packets * sizeof(AVPacket *));
+
+  for (int i = 0; i < encoder->max_num_packets; i++) {
+    encoder->packets[i] = av_packet_alloc();
+  }
+
+  return encoder;
+}
+
+int encoder_init(struct Encoder *encoder, struct EncoderConfig *config) {
+  encoder->codec = avcodec_find_encoder(config->codec);
+  if (!encoder->codec) {
+    return -1;
+  }
+
+  encoder->c = avcodec_alloc_context3(encoder->codec);
+  if (!encoder->c) {
+    return -1;
+  }
+
+  encoder->c->width = config->width;
+  encoder->c->height = config->height;
+  encoder->c->pix_fmt = config->format;
+  encoder->c->time_base = config->time_base;
+
+  if (config->profile != FF_PROFILE_UNKNOWN) {
+    encoder->c->profile = config->profile;
+  }
+
+  if (config->gop_size > 0) {
+    encoder->c->gop_size = config->gop_size;
+  }
+
+  if (config->max_b_frames >= 0) {
+    encoder->c->max_b_frames = config->max_b_frames;
+  }
+
+  AVDictionary *opts = NULL;
+  if (config->codec == AV_CODEC_ID_HEVC) {
+    char x265_params[256] = "log-level=warning";
+    if (config->gop_size > 0) {
+      sprintf(x265_params + strlen(x265_params), ":keyint=%d", config->gop_size);
+    }
+
+    if (config->max_b_frames >= 0) {
+      sprintf(x265_params + strlen(x265_params), ":bframes=%d", config->max_b_frames);
+    }
+
+    av_dict_set(&opts, "x265-params", x265_params, 0);
+  }
+
+  return avcodec_open2(encoder->c, encoder->codec, &opts);
+}
+
+int encoder_encode(struct Encoder *encoder, AVFrame *frame) {
+  int ret = avcodec_send_frame(encoder->c, frame);
+  if (ret < 0) {
+    return ret;
+  }
+
+  encoder->num_packets = 0;
+
+  while (1) {
+    ret = avcodec_receive_packet(encoder->c, encoder->packets[encoder->num_packets]);
+    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
+      break;
+    } else if (ret < 0) {
+      return ret;
+    }
+
+    if (++encoder->num_packets >= encoder->max_num_packets) {
+      encoder->max_num_packets *= 2;
+      encoder->packets =
+          XAV_REALLOC(encoder->packets, encoder->max_num_packets * sizeof(AVPacket *));
+      for (int i = encoder->num_packets; i < encoder->max_num_packets; i++) {
+        encoder->packets[i] = av_packet_alloc();
+      }
+    }
+  }
+
+  return 0;
+}
+
+void encoder_free(struct Encoder **encoder) {
+  if (*encoder != NULL) {
+    struct Encoder *e = *encoder;
+
+    if (e->c != NULL) {
+      avcodec_free_context(&e->c);
+    }
+
+    for (int i = 0; i < e->max_num_packets; i++) {
+      av_packet_free(&e->packets[i]);
+    }
+
+    XAV_FREE(e);
+    *encoder = NULL;
+  }
+}

c_src/xav/encoder.h

@@ -0,0 +1,30 @@
+#include "utils.h"
+#include <libavcodec/avcodec.h>
+
+struct Encoder {
+  const AVCodec *codec;
+  AVCodecContext *c;
+  int num_packets;
+  int max_num_packets;
+  AVPacket **packets;
+};
+
+struct EncoderConfig {
+  enum AVMediaType media_type;
+  enum AVCodecID codec;
+  int width;
+  int height;
+  enum AVPixelFormat format;
+  AVRational time_base;
+  int gop_size;
+  int max_b_frames;
+  int profile;
+};
+
+struct Encoder *encoder_alloc();
+
+int encoder_init(struct Encoder *encoder, struct EncoderConfig *encoder_config);
+
+int encoder_encode(struct Encoder *encoder, AVFrame *frame);
+
+void encoder_free(struct Encoder **encoder);

c_src/xav/utils.c

@@ -1,5 +1,4 @@
 #include "utils.h"
-#include <libavutil/imgutils.h>
 #include <libavutil/mathematics.h>
 #include <libavutil/opt.h>
 #include <stdint.h>
@@ -20,19 +19,19 @@ ERL_NIF_TERM xav_nif_raise(ErlNifEnv *env, char *msg) {
   return enif_raise_exception(env, reason);
 }
 
-int xav_get_atom(ErlNifEnv *env, ERL_NIF_TERM atom, char **value) {
+int xav_nif_get_atom(ErlNifEnv *env, ERL_NIF_TERM term, char **value) {
   unsigned int atom_len;
-  if (!enif_get_atom_length(env, atom, &atom_len, ERL_NIF_LATIN1)) {
+  if (!enif_get_atom_length(env, term, &atom_len, ERL_NIF_LATIN1)) {
     return 0;
   }
 
-  char *format = (char *)XAV_ALLOC((atom_len * 1) * sizeof(char *));
-  if (!enif_get_atom(env, atom, format, atom_len + 1, ERL_NIF_LATIN1)) {
-    XAV_FREE(format);
+  char *atom_value = (char *)XAV_ALLOC((atom_len + 1) * sizeof(char *));
+  if (!enif_get_atom(env, term, atom_value, atom_len + 1, ERL_NIF_LATIN1)) {
+    XAV_FREE(atom_value);
     return 0;
   }
 
-  *value = format;
+  *value = atom_value;
   return 1;
 }
 
@@ -66,3 +65,17 @@ ERL_NIF_TERM xav_nif_video_frame_to_term(ErlNifEnv *env, AVFrame *frame) {
   ERL_NIF_TERM pts_term = enif_make_int64(env, frame->pts);
   return enif_make_tuple(env, 5, data_term, format_term, width_term, height_term, pts_term);
 }
+
+ERL_NIF_TERM xav_nif_packet_to_term(ErlNifEnv *env, AVPacket *packet) {
+  ERL_NIF_TERM data_term;
+
+  unsigned char *ptr = enif_make_new_binary(env, packet->size, &data_term);
+
+  memcpy(ptr, packet->data, packet->size);
+
+  ERL_NIF_TERM dts = enif_make_int(env, packet->dts);
+  ERL_NIF_TERM pts = enif_make_int(env, packet->pts);
+  ERL_NIF_TERM is_keyframe =
+      enif_make_atom(env, packet->flags & AV_PKT_FLAG_KEY ? "true" : "false");
+  return enif_make_tuple(env, 4, data_term, dts, pts, is_keyframe);
+}

c_src/xav/utils.h

@@ -3,6 +3,7 @@
 #include <libavdevice/avdevice.h>
 #include <libavformat/avformat.h>
 #include <libavutil/avutil.h>
+#include <libavutil/imgutils.h>
 #include <libswresample/swresample.h>
 #include <libswscale/swscale.h>
 
@@ -14,11 +15,13 @@
 #endif
 
 #define XAV_ALLOC(X) enif_alloc(X)
+#define XAV_REALLOC(X, Y) enif_realloc(X, Y)
 #define XAV_FREE(X) enif_free(X)
 ERL_NIF_TERM xav_nif_ok(ErlNifEnv *env, ERL_NIF_TERM data_term);
 ERL_NIF_TERM xav_nif_error(ErlNifEnv *env, char *reason);
 ERL_NIF_TERM xav_nif_raise(ErlNifEnv *env, char *msg);
-int xav_get_atom(ErlNifEnv *env, ERL_NIF_TERM atom, char **value);
+int xav_nif_get_atom(ErlNifEnv *env, ERL_NIF_TERM term, char **value);
 ERL_NIF_TERM xav_nif_video_frame_to_term(ErlNifEnv *env, AVFrame *frame);
 ERL_NIF_TERM xav_nif_audio_frame_to_term(ErlNifEnv *env, uint8_t **out_data, int out_samples,
                                          int out_size, enum AVSampleFormat out_format, int pts);
+ERL_NIF_TERM xav_nif_packet_to_term(ErlNifEnv *env, AVPacket *packet);

c_src/xav/xav_decoder.c

@@ -22,7 +22,7 @@ ERL_NIF_TERM new (ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
   char *out_format = NULL;
 
   // resolve codec
-  if (!xav_get_atom(env, argv[0], &codec)) {
+  if (!xav_nif_get_atom(env, argv[0], &codec)) {
     return xav_nif_raise(env, "failed_to_get_atom");
   }
 
@@ -46,7 +46,7 @@ ERL_NIF_TERM new (ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
   }
 
   // resolve output format
-  if (!xav_get_atom(env, argv[1], &out_format)) {
+  if (!xav_nif_get_atom(env, argv[1], &out_format)) {
     ret = xav_nif_raise(env, "failed_to_get_atom");
     goto clean;
   }