Allow for specifying target bitrate in options

c_src/membrane_vpx_plugin/vpx_common.c

@@ -73,13 +73,3 @@ void convert_between_image_and_raw_frame(
     }
   }
 }
-
-void free_payloads(UnifexPayload **payloads, unsigned int payloads_cnt) {
-  for (unsigned int i = 0; i < payloads_cnt; i++) {
-    if (payloads[i] != NULL) {
-      unifex_payload_release(payloads[i]);
-      unifex_free(payloads[i]);
-    }
-  }
-  unifex_free(payloads);
-}

c_src/membrane_vpx_plugin/vpx_common.h

@@ -1,8 +1,8 @@
 #pragma once
+#include "unifex/payload.h"
+#include "unifex/unifex.h"
 #include "vpx/vpx_codec.h"
 #include "vpx/vpx_image.h"
-#include <unifex/payload.h>
-#include <unifex/unifex.h>
 
 typedef struct Dimensions {
   unsigned int width;
@@ -21,8 +21,6 @@ typedef enum ConversionType { IMAGE_TO_RAW_FRAME, RAW_FRAME_TO_IMAGE } Conversio
 
 Dimensions get_plane_dimensions(const vpx_image_t *img, int plane);
 
-void free_payloads(UnifexPayload **payloads, unsigned int payloads_cnt);
-
 void convert_between_image_and_raw_frame(
     vpx_image_t *img, UnifexPayload *raw_frame, ConversionType conversion_type
-);
+);

c_src/membrane_vpx_plugin/vpx_decoder.c

@@ -52,6 +52,16 @@ void alloc_output_frame(UnifexEnv *env, const vpx_image_t *img, UnifexPayload **
   unifex_payload_alloc(env, UNIFEX_PAYLOAD_BINARY, get_image_byte_size(img), *output_frame);
 }
 
+void free_payloads(UnifexPayload **payloads, unsigned int payloads_cnt) {
+  for (unsigned int i = 0; i < payloads_cnt; i++) {
+    if (payloads[i] != NULL) {
+      unifex_payload_release(payloads[i]);
+      unifex_free(payloads[i]);
+    }
+  }
+  unifex_free(payloads);
+}
+
 PixelFormat get_pixel_format_from_image(vpx_image_t *img) {
   switch (img->fmt) {
   case VPX_IMG_FMT_I422:
@@ -74,7 +84,7 @@ UNIFEX_TERM decode_frame(UnifexEnv *env, UnifexPayload *frame, State *state) {
   vpx_image_t *img = NULL;
   PixelFormat pixel_format = PIXEL_FORMAT_I420;
   unsigned int frames_cnt = 0, allocated_frames = 1;
-  UnifexPayload **output_frames = unifex_alloc(allocated_frames * sizeof(UnifexPayload*));
+  UnifexPayload **output_frames = unifex_alloc(allocated_frames * sizeof(UnifexPayload *));
 
   if (vpx_codec_decode(&state->codec_context, frame->data, frame->size, NULL, 0)) {
     return result_error(

c_src/membrane_vpx_plugin/vpx_encoder.c

@@ -1,4 +1,7 @@
 #include "vpx_encoder.h"
+#include "membrane_vpx_plugin/_generated/nif/vpx_encoder.h"
+#include "unifex/payload.h"
+#include <stdio.h>
 
 // The following code is based on the simple_encoder example provided by libvpx
 // (https://github.com/webmproject/libvpx/blob/main/examples/simple_encoder.c)
@@ -26,14 +29,7 @@ vpx_img_fmt_t translate_pixel_format(PixelFormat pixel_format) {
   }
 }
 
-UNIFEX_TERM create(
-    UnifexEnv *env,
-    Codec codec,
-    unsigned int width,
-    unsigned int height,
-    PixelFormat pixel_format,
-    unsigned int encoding_deadline
-) {
+UNIFEX_TERM create(UnifexEnv *env, Codec codec, encoder_options opts) {
   UNIFEX_TERM result;
   State *state = unifex_alloc_state(env);
   vpx_codec_enc_cfg_t config;
@@ -46,24 +42,27 @@ UNIFEX_TERM create(
     state->codec_interface = vpx_codec_vp9_cx();
     break;
   }
-  state->encoding_deadline = encoding_deadline;
+  state->encoding_deadline = opts.encoding_deadline;
 
   if (vpx_codec_enc_config_default(state->codec_interface, &config, 0)) {
     return result_error(
         env, "Failed to get default codec config", create_result_error, NULL, state
     );
   }
 
-  config.g_h = height;
-  config.g_w = width;
+  config.g_h = opts.height;
+  config.g_w = opts.width;
+  config.rc_target_bitrate = opts.target_bitrate;
   config.g_timebase.num = 1;
   config.g_timebase.den = 1000000000; // 1e9
   config.g_error_resilient = 1;
 
   if (vpx_codec_enc_init(&state->codec_context, state->codec_interface, &config, 0)) {
     return result_error(env, "Failed to initialize encoder", create_result_error, NULL, state);
   }
-  if (!vpx_img_alloc(&state->img, translate_pixel_format(pixel_format), width, height, 1)) {
+  if (!vpx_img_alloc(
+          &state->img, translate_pixel_format(opts.pixel_format), opts.width, opts.height, 1
+      )) {
     return result_error(
         env, "Failed to allocate image", create_result_error, &state->codec_context, state
     );
@@ -78,30 +77,44 @@ void get_image_from_raw_frame(vpx_image_t *img, UnifexPayload *raw_frame) {
 }
 
 void alloc_output_frame(
-    UnifexEnv *env, const vpx_codec_cx_pkt_t *packet, UnifexPayload **output_frame
+    UnifexEnv *env, const vpx_codec_cx_pkt_t *packet, encoded_frame *output_frame
 ) {
-  *output_frame = unifex_alloc(sizeof(UnifexPayload));
-  unifex_payload_alloc(env, UNIFEX_PAYLOAD_BINARY, packet->data.frame.sz, *output_frame);
+  output_frame->payload = unifex_alloc(sizeof(UnifexPayload));
+  unifex_payload_alloc(env, UNIFEX_PAYLOAD_BINARY, packet->data.frame.sz, output_frame->payload);
+}
+
+void free_frames(encoded_frame *frames, unsigned int frames_cnt) {
+  for (unsigned int i = 0; i < frames_cnt; i++) {
+    UnifexPayload *payload = frames[i].payload;
+    if (payload != NULL) {
+      unifex_payload_release(payload);
+      unifex_free(payload);
+    }
+  }
+  unifex_free(frames);
 }
 
-UNIFEX_TERM encode(UnifexEnv *env, vpx_image_t *img, vpx_codec_pts_t pts, State *state) {
+UNIFEX_TERM encode(
+    UnifexEnv *env, vpx_image_t *img, vpx_codec_pts_t pts, int force_keyframe, State *state
+) {
   vpx_codec_iter_t iter = NULL;
   int flushing = (img == NULL), got_packets = 0;
   const vpx_codec_cx_pkt_t *packet = NULL;
 
   unsigned int frames_cnt = 0, allocated_frames = 1;
-  UnifexPayload **encoded_frames = unifex_alloc(allocated_frames * sizeof(UnifexPayload*));
-  vpx_codec_pts_t *encoded_frames_timestamps =
-      unifex_alloc(allocated_frames * sizeof(vpx_codec_pts_t));
+  encoded_frame *encoded_frames = unifex_alloc(allocated_frames * sizeof(encoded_frame));
 
   do {
-    // Reasoning for the do-while and while loops comes from the description of vpx_codec_encode:
+    // Reasoning for the do-while and while loops comes from the description of
+    // vpx_codec_encode:
     //
-    // When the last frame has been passed to the encoder, this function should continue to be
-    // called, with the img parameter set to NULL. This will signal the end-of-stream condition to
-    // the encoder and allow it to encode any held buffers. Encoding is complete when
-    // vpx_codec_encode() is called and vpx_codec_get_cx_data() returns no data.
-    if (vpx_codec_encode(&state->codec_context, img, pts, 1, 0, state->encoding_deadline) !=
+    // When the last frame has been passed to the encoder, this function should
+    // continue to be called, with the img parameter set to NULL. This will
+    // signal the end-of-stream condition to the encoder and allow it to encode
+    // any held buffers. Encoding is complete when vpx_codec_encode() is called
+    // and vpx_codec_get_cx_data() returns no data.
+    vpx_enc_frame_flags_t flags = force_keyframe ? VPX_EFLAG_FORCE_KF : 0;
+    if (vpx_codec_encode(&state->codec_context, img, pts, 1, flags, state->encoding_deadline) !=
         VPX_CODEC_OK) {
       if (flushing) {
         return result_error(
@@ -113,6 +126,7 @@ UNIFEX_TERM encode(UnifexEnv *env, vpx_image_t *img, vpx_codec_pts_t pts, State
         );
       }
     }
+    force_keyframe = 0;
     got_packets = 0;
 
     while ((packet = vpx_codec_get_cx_data(&state->codec_context, &iter)) != NULL) {
@@ -121,38 +135,36 @@ UNIFEX_TERM encode(UnifexEnv *env, vpx_image_t *img, vpx_codec_pts_t pts, State
 
       if (frames_cnt >= allocated_frames) {
         allocated_frames *= 2;
-        encoded_frames = unifex_realloc(encoded_frames, allocated_frames * sizeof(*encoded_frames));
-
-        encoded_frames_timestamps = unifex_realloc(
-            encoded_frames_timestamps, allocated_frames * sizeof(*encoded_frames_timestamps)
-        );
+        encoded_frames = unifex_realloc(encoded_frames, allocated_frames * sizeof(encoded_frame));
       }
       alloc_output_frame(env, packet, &encoded_frames[frames_cnt]);
-      memcpy(encoded_frames[frames_cnt]->data, packet->data.frame.buf, packet->data.frame.sz);
-      encoded_frames_timestamps[frames_cnt] = packet->data.frame.pts;
+      memcpy(
+          encoded_frames[frames_cnt].payload->data, packet->data.frame.buf, packet->data.frame.sz
+      );
+      encoded_frames[frames_cnt].pts = packet->data.frame.pts;
+      encoded_frames[frames_cnt].is_keyframe = ((packet->data.frame.flags & VPX_FRAME_IS_KEY) != 0);
       frames_cnt++;
     }
   } while (got_packets && flushing);
 
   UNIFEX_TERM result;
   if (flushing) {
-    result =
-        flush_result_ok(env, encoded_frames, frames_cnt, encoded_frames_timestamps, frames_cnt);
+    result = flush_result_ok(env, encoded_frames, frames_cnt);
   } else {
-    result = encode_frame_result_ok(
-        env, encoded_frames, frames_cnt, encoded_frames_timestamps, frames_cnt
-    );
+    result = encode_frame_result_ok(env, encoded_frames, frames_cnt);
   }
-  free_payloads(encoded_frames, frames_cnt);
+  free_frames(encoded_frames, frames_cnt);
 
   return result;
 }
 
 UNIFEX_TERM encode_frame(
-    UnifexEnv *env, UnifexPayload *raw_frame, vpx_codec_pts_t pts, State *state
+    UnifexEnv *env, UnifexPayload *raw_frame, vpx_codec_pts_t pts, int force_keyframe, State *state
 ) {
   get_image_from_raw_frame(&state->img, raw_frame);
-  return encode(env, &state->img, pts, state);
+  return encode(env, &state->img, pts, force_keyframe, state);
 }
 
-UNIFEX_TERM flush(UnifexEnv *env, State *state) { return encode(env, NULL, 0, state); }
+UNIFEX_TERM flush(UnifexEnv *env, int force_keyframe, State *state) {
+  return encode(env, NULL, 0, force_keyframe, state);
+}

c_src/membrane_vpx_plugin/vpx_encoder.spec.exs

@@ -6,21 +6,29 @@ type codec :: :vp8 | :vp9
 
 type pixel_format :: :I420 | :I422 | :I444 | :NV12 | :YV12
 
-spec create(
-       codec,
-       width :: unsigned,
-       height :: unsigned,
-       pixel_format,
-       encoding_deadline :: unsigned
-     ) ::
+type encoded_frame :: %EncodedFrame{
+       payload: payload,
+       pts: int64,
+       is_keyframe: bool
+     }
+
+type encoder_options :: %EncoderOptions{
+       width: unsigned,
+       height: unsigned,
+       pixel_format: pixel_format,
+       encoding_deadline: unsigned,
+       target_bitrate: unsigned
+     }
+
+spec create(codec, encoder_options) ::
        {:ok :: label, state} | {:error :: label, reason :: atom}
 
-spec encode_frame(payload, pts :: int64, state) ::
-       {:ok :: label, frames :: [payload], timestamps :: [int64]}
+spec encode_frame(payload, pts :: int64, force_keyframe :: bool, state) ::
+       {:ok :: label, frames :: [encoded_frame]}
        | {:error :: label, reason :: atom}
 
-spec flush(state) ::
-       {:ok :: label, frames :: [payload], timestamps :: [int64]}
+spec flush(force_keyframe :: bool, state) ::
+       {:ok :: label, frames :: [encoded_frame]}
        | {:error :: label, reason :: atom}
 
 dirty :cpu, [:create, :encode_frame, :flush]

lib/membrane_vpx/encoder/vp8_encoder.ex

@@ -1,6 +1,15 @@
 defmodule Membrane.VP8.Encoder do
   @moduledoc """
-  Element that encodes a VP8 stream
+  Element that encodes a VP8 stream.
+
+  This element can receive a `Membrane.VPx.ForceKeyframeEvent` on it's `:output` pad to force the
+  next frame to be a keyframe.
+
+  Buffers produced by this element will have the following metadata that inform whether the buffer
+  contains a keyframe:
+  ```elixir
+  %{vp8: %{is_keyframe: is_keyframe :: boolean()}}
+  ```
   """
   use Membrane.Filter
 
@@ -17,6 +26,13 @@ defmodule Membrane.VP8.Encoder do
                 If set to `:auto` the deadline will be calculated based on the framerate provided by
                 incoming stream format. If the framerate is `nil` a fixed deadline of 10ms will be set.
                 """
+              ],
+              target_bitrate: [
+                spec: pos_integer(),
+                default: 1000,
+                description: """
+                Gives the encoder information about the target bitrate (in kb/s).
+                """
               ]
 
   def_input_pad :input,
@@ -36,6 +52,9 @@ defmodule Membrane.VP8.Encoder do
   @impl true
   defdelegate handle_buffer(pad, buffer, ctx, state), to: VPx.Encoder
 
+  @impl true
+  defdelegate handle_event(pad, event, ctx, state), to: VPx.Encoder
+
   @impl true
   defdelegate handle_end_of_stream(pad, ctx, state), to: VPx.Encoder
 end

lib/membrane_vpx/encoder/vp9_encoder.ex

@@ -1,6 +1,15 @@
 defmodule Membrane.VP9.Encoder do
   @moduledoc """
-  Element that encodes a VP9 stream
+  Element that encodes a VP9 stream.
+
+  This element can receive a `Membrane.VPx.ForceKeyframeEvent` on it's `:output` pad to force the
+  next frame to be a keyframe.
+
+  Buffers produced by this element will have the following metadata that inform whether the buffer
+  contains a keyframe:
+  ```elixir
+  %{vp9: %{is_keyframe: is_keyframe :: boolean()}}
+  ```
   """
   use Membrane.Filter
 
@@ -17,6 +26,13 @@ defmodule Membrane.VP9.Encoder do
                 If set to `:auto` the deadline will be calculated based on the framerate provided by
                 incoming stream format. If the framerate is `nil` a fixed deadline of 10ms will be set.
                 """
+              ],
+              target_bitrate: [
+                spec: pos_integer(),
+                default: 1000,
+                description: """
+                Gives the encoder information about the target bitrate (in kb/s).
+                """
               ]
 
   def_input_pad :input,
@@ -36,6 +52,9 @@ defmodule Membrane.VP9.Encoder do
   @impl true
   defdelegate handle_buffer(pad, buffer, ctx, state), to: VPx.Encoder
 
+  @impl true
+  defdelegate handle_event(pad, event, ctx, state), to: VPx.Encoder
+
   @impl true
   defdelegate handle_end_of_stream(pad, ctx, state), to: VPx.Encoder
 end