Add encoder_options

membraneframework · Jul 11, 2024 · 51a190c · 51a190c
1 parent 413d2f1
commit 51a190c
Show file tree

Hide file tree

Showing 7 changed files with 57 additions and 34 deletions.
diff --git a/c_src/membrane_vpx_plugin/vpx_encoder.c b/c_src/membrane_vpx_plugin/vpx_encoder.c
@@ -29,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;
@@ -49,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
     );

diff --git a/c_src/membrane_vpx_plugin/vpx_encoder.spec.exs b/c_src/membrane_vpx_plugin/vpx_encoder.spec.exs
@@ -12,13 +12,15 @@ type encoded_frame :: %EncodedFrame{
        is_keyframe: bool
      }
 
-spec create(
-       codec,
-       width :: unsigned,
-       height :: unsigned,
-       pixel_format,
-       encoding_deadline :: unsigned
-     ) ::
+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, force_keyframe :: bool, state) ::

diff --git a/lib/membrane_vpx/encoder/vp8_encoder.ex b/lib/membrane_vpx/encoder/vp8_encoder.ex
@@ -26,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,

diff --git a/lib/membrane_vpx/encoder/vp9_encoder.ex b/lib/membrane_vpx/encoder/vp9_encoder.ex
@@ -26,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,

diff --git a/lib/membrane_vpx/encoder/vpx_encoder.ex b/lib/membrane_vpx/encoder/vpx_encoder.ex
@@ -14,11 +14,12 @@ defmodule Membrane.VPx.Encoder do
             codec: :vp8 | :vp9,
             codec_module: VP8 | VP9,
             encoding_deadline: non_neg_integer(),
+            target_bitrate: pos_integer(),
             encoder_ref: reference() | nil,
             force_next_keyframe: boolean()
           }
 
-    @enforce_keys [:codec, :codec_module, :encoding_deadline]
+    @enforce_keys [:codec, :codec_module, :encoding_deadline, :target_bitrate]
     defstruct @enforce_keys ++
                 [
                   encoder_ref: nil,
@@ -28,6 +29,14 @@ defmodule Membrane.VPx.Encoder do
 
   @type callback_return :: {[Membrane.Element.Action.t()], State.t()}
 
+  @type encoder_options :: %{
+          width: pos_integer(),
+          height: pos_integer(),
+          pixel_format: Membrane.RawVideo.pixel_format(),
+          encoding_deadline: non_neg_integer(),
+          target_bitrate: pos_integer()
+        }
+
   @spec handle_init(CallbackContext.t(), VP8.Encoder.t() | VP9.Encoder.t(), :vp8 | :vp9) ::
           callback_return()
 
@@ -39,7 +48,8 @@ defmodule Membrane.VPx.Encoder do
           :vp8 -> VP8
           :vp9 -> VP9
         end,
-      encoding_deadline: opts.encoding_deadline
+      encoding_deadline: opts.encoding_deadline,
+      target_bitrate: opts.target_bitrate
     }
 
     {[], state}
@@ -115,8 +125,15 @@ defmodule Membrane.VPx.Encoder do
         {fixed_deadline, _framerate} -> fixed_deadline |> Membrane.Time.as_microseconds(:round)
       end
 
-    new_encoder_ref =
-      Native.create!(state.codec, width, height, pixel_format, encoding_deadline)
+    encoder_options = %{
+      width: width,
+      height: height,
+      pixel_format: pixel_format,
+      encoding_deadline: encoding_deadline,
+      target_bitrate: state.target_bitrate
+    }
+
+    new_encoder_ref = Native.create!(state.codec, encoder_options)
 
     case state.encoder_ref do
       nil ->

diff --git a/lib/membrane_vpx/encoder/vpx_encoder_native.ex b/lib/membrane_vpx/encoder/vpx_encoder_native.ex
@@ -2,15 +2,9 @@ defmodule Membrane.VPx.Encoder.Native do
   @moduledoc false
   use Unifex.Loader
 
-  @spec create!(
-          :vp8 | :vp9,
-          pos_integer(),
-          pos_integer(),
-          Membrane.RawVideo.pixel_format(),
-          non_neg_integer()
-        ) :: reference()
-  def create!(codec, width, height, pixel_format, encoding_deadline) do
-    case create(codec, width, height, pixel_format, encoding_deadline) do
+  @spec create!(:vp8 | :vp9, Membrane.VPx.Encoder.encoder_options()) :: reference()
+  def create!(codec, encoder_options) do
+    case create(codec, encoder_options) do
       {:ok, decoder_ref} -> decoder_ref
       {:error, reason} -> raise "Failed to create native encoder: #{inspect(reason)}"
     end

diff --git a/test/membrane_vpx_plugin/encoder/vpx_encoder_test.exs b/test/membrane_vpx_plugin/encoder/vpx_encoder_test.exs
@@ -14,7 +14,7 @@ defmodule Membrane.VPx.EncoderTest do
         "ref_vp8.raw",
         "output_vp8.ivf",
         "ref_vp8.ivf",
-        %Membrane.VP8.Encoder{encoding_deadline: 0}
+        %Membrane.VP8.Encoder{encoding_deadline: 0, target_bitrate: 256}
       )
     end
 
@@ -24,7 +24,7 @@ defmodule Membrane.VPx.EncoderTest do
         "ref_vp9.raw",
         "output_vp9.ivf",
         "ref_vp9.ivf",
-        %Membrane.VP9.Encoder{encoding_deadline: 0}
+        %Membrane.VP9.Encoder{encoding_deadline: 0, target_bitrate: 256}
       )
     end
   end