Lag in frames option (#7)

* Finish implementing lag in frames * Improve keyframes test * Add ZeroLatencyTest
membraneframework · Aug 19, 2024 · 3363507 · 3363507
1 parent 8b6800d
commit 3363507
Show file tree

Hide file tree

Showing 8 changed files with 211 additions and 48 deletions.
diff --git a/c_src/membrane_vpx_plugin/vpx_encoder.c b/c_src/membrane_vpx_plugin/vpx_encoder.c
@@ -1,7 +1,6 @@
 #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)
@@ -29,7 +28,20 @@ vpx_img_fmt_t translate_pixel_format(PixelFormat pixel_format) {
   }
 }
 
-UNIFEX_TERM create(UnifexEnv *env, Codec codec, encoder_options opts) {
+void apply_user_encoder_config(vpx_codec_enc_cfg_t *config, user_encoder_config *user_config) {
+  config->g_lag_in_frames = user_config->g_lag_in_frames;
+  config->rc_target_bitrate = user_config->rc_target_bitrate;
+}
+
+UNIFEX_TERM create(
+    UnifexEnv *env,
+    Codec codec,
+    unsigned int width,
+    unsigned int height,
+    PixelFormat pixel_format,
+    unsigned int encoding_deadline,
+    user_encoder_config user_config
+) {
   UNIFEX_TERM result;
   State *state = unifex_alloc_state(env);
   vpx_codec_enc_cfg_t config;
@@ -42,27 +54,25 @@ UNIFEX_TERM create(UnifexEnv *env, Codec codec, encoder_options opts) {
     state->codec_interface = vpx_codec_vp9_cx();
     break;
   }
-  state->encoding_deadline = opts.encoding_deadline;
+  state->encoding_deadline = 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 = opts.height;
-  config.g_w = opts.width;
-  config.rc_target_bitrate = opts.rc_target_bitrate;
+  config.g_h = height;
+  config.g_w = width;
   config.g_timebase.num = 1;
   config.g_timebase.den = 1000000000; // 1e9
   config.g_error_resilient = 1;
+  apply_user_encoder_config(&config, &user_config);
 
   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(opts.pixel_format), opts.width, opts.height, 1
-      )) {
+  if (!vpx_img_alloc(&state->img, translate_pixel_format(pixel_format), width, 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,16 +12,19 @@ type encoded_frame :: %EncodedFrame{
        is_keyframe: bool
      }
 
-type encoder_options :: %EncoderOptions{
-       width: unsigned,
-       height: unsigned,
-       pixel_format: pixel_format,
-       encoding_deadline: unsigned,
+type user_encoder_config :: %UserEncoderConfig{
+       g_lag_in_frames: unsigned,
        rc_target_bitrate: unsigned
      }
 
-spec create(codec, encoder_options) ::
-       {:ok :: label, state} | {:error :: label, reason :: atom}
+spec create(
+       codec,
+       width :: unsigned,
+       height :: unsigned,
+       pixel_format,
+       encoding_deadline :: unsigned,
+       user_encoder_config
+     ) :: {:ok :: label, state} | {:error :: label, reason :: atom}
 
 spec encode_frame(payload, pts :: int64, force_keyframe :: bool, state) ::
        {:ok :: label, frames :: [encoded_frame]}

diff --git a/lib/membrane_vpx/encoder/vp8_encoder.ex b/lib/membrane_vpx/encoder/vp8_encoder.ex
@@ -36,6 +36,20 @@ defmodule Membrane.VP8.Encoder do
                 of the incoming stream. Some reference recommended bitrates can be also found
                 [here](https://support.google.com/youtube/answer/1722171#zippy=%2Cbitrate)
                 """
+              ],
+              g_lag_in_frames: [
+                spec: non_neg_integer(),
+                default: 5,
+                description: """
+                The number of input frames the encoder is allowed to consume
+                before producing output frames. This allows the encoder to
+                base decisions for the current frame on future frames. This does
+                increase the latency of the encoding pipeline, so it is not appropriate
+                in all situations (ex: realtime encoding).
+
+                Note that this is a maximum value -- the encoder may produce frames
+                sooner than the given limit. If set to 0 this feature will be disabled.
+                """
               ]
 
   def_input_pad :input,

diff --git a/lib/membrane_vpx/encoder/vp9_encoder.ex b/lib/membrane_vpx/encoder/vp9_encoder.ex
@@ -36,6 +36,20 @@ defmodule Membrane.VP9.Encoder do
                 of the incoming stream. Some reference recommended bitrates can be also found
                 [here](https://support.google.com/youtube/answer/1722171#zippy=%2Cbitrate)
                 """
+              ],
+              g_lag_in_frames: [
+                spec: non_neg_integer(),
+                default: 5,
+                description: """
+                The number of input frames the encoder is allowed to consume
+                before producing output frames. This allows the encoder to
+                base decisions for the current frame on future frames. This does
+                increase the latency of the encoding pipeline, so it is not appropriate
+                in all situations (ex: realtime encoding).
+
+                Note that this is a maximum value -- the encoder may produce frames
+                sooner than the given limit. If set to 0 this feature will be disabled.
+                """
               ]
 
   def_input_pad :input,

diff --git a/lib/membrane_vpx/encoder/vpx_encoder.ex b/lib/membrane_vpx/encoder/vpx_encoder.ex
@@ -8,19 +8,28 @@ defmodule Membrane.VPx.Encoder do
   @default_encoding_deadline Membrane.Time.milliseconds(10)
   @bitrate_calculation_coefficient 0.14
 
+  @type unprocessed_user_encoder_config :: %{
+          g_lag_in_frames: non_neg_integer(),
+          rc_target_bitrate: pos_integer() | :auto
+        }
+  @type user_encoder_config :: %{
+          g_lag_in_frames: non_neg_integer(),
+          rc_target_bitrate: pos_integer()
+        }
+
   defmodule State do
     @moduledoc false
 
     @type t :: %__MODULE__{
             codec: :vp8 | :vp9,
             codec_module: VP8 | VP9,
-            encoding_deadline: non_neg_integer(),
-            rc_target_bitrate: pos_integer(),
+            encoding_deadline: non_neg_integer() | :auto,
+            user_encoder_config: Membrane.VPx.Encoder.unprocessed_user_encoder_config(),
             encoder_ref: reference() | nil,
             force_next_keyframe: boolean()
           }
 
-    @enforce_keys [:codec, :codec_module, :encoding_deadline, :rc_target_bitrate]
+    @enforce_keys [:codec, :codec_module, :encoding_deadline, :user_encoder_config]
     defstruct @enforce_keys ++
                 [
                   encoder_ref: nil,
@@ -30,14 +39,6 @@ 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(),
-          rc_target_bitrate: pos_integer()
-        }
-
   @type encoded_frame :: %{payload: binary(), pts: non_neg_integer(), is_keyframe: boolean()}
 
   @spec handle_init(CallbackContext.t(), VP8.Encoder.t() | VP9.Encoder.t(), :vp8 | :vp9) ::
@@ -52,7 +53,10 @@ defmodule Membrane.VPx.Encoder do
           :vp9 -> VP9
         end,
       encoding_deadline: opts.encoding_deadline,
-      rc_target_bitrate: opts.rc_target_bitrate
+      user_encoder_config: %{
+        g_lag_in_frames: opts.g_lag_in_frames,
+        rc_target_bitrate: opts.rc_target_bitrate
+      }
     }
 
     {[], state}
@@ -128,17 +132,18 @@ defmodule Membrane.VPx.Encoder do
         {fixed_deadline, _framerate} -> fixed_deadline |> Membrane.Time.as_microseconds(:round)
       end
 
-    rc_target_bitrate = get_target_bitrate(state.rc_target_bitrate, width, height, framerate)
+    user_encoder_config =
+      process_user_encoder_config(state.user_encoder_config, width, height, framerate)
 
-    encoder_options = %{
-      width: width,
-      height: height,
-      pixel_format: pixel_format,
-      encoding_deadline: encoding_deadline,
-      rc_target_bitrate: rc_target_bitrate
-    }
-
-    new_encoder_ref = Native.create!(state.codec, encoder_options)
+    new_encoder_ref =
+      Native.create!(
+        state.codec,
+        width,
+        height,
+        pixel_format,
+        encoding_deadline,
+        user_encoder_config
+      )
 
     case state.encoder_ref do
       nil ->
@@ -149,13 +154,29 @@ defmodule Membrane.VPx.Encoder do
     end
   end
 
-  @spec get_target_bitrate(
+  @spec process_user_encoder_config(
+          unprocessed_user_encoder_config(),
+          pos_integer(),
+          pos_integer(),
+          {non_neg_integer(), pos_integer()} | nil
+        ) :: user_encoder_config()
+  defp process_user_encoder_config(user_encoder_config, width, height, framerate) do
+    rc_target_bitrate =
+      process_rc_target_bitrate(user_encoder_config.rc_target_bitrate, width, height, framerate)
+
+    %{
+      g_lag_in_frames: user_encoder_config.g_lag_in_frames,
+      rc_target_bitrate: rc_target_bitrate
+    }
+  end
+
+  @spec process_rc_target_bitrate(
           pos_integer() | :auto,
           pos_integer(),
           pos_integer(),
           {non_neg_integer(), pos_integer()} | nil
         ) :: pos_integer()
-  defp get_target_bitrate(:auto, width, height, framerate) do
+  defp process_rc_target_bitrate(:auto, width, height, framerate) do
     assumed_fps =
       case framerate do
         nil -> 30.0
@@ -165,7 +186,7 @@ defmodule Membrane.VPx.Encoder do
     (@bitrate_calculation_coefficient * width * height * assumed_fps) |> trunc() |> div(1000)
   end
 
-  defp get_target_bitrate(provided_bitrate, _width, _height, _framerate) do
+  defp process_rc_target_bitrate(provided_bitrate, _width, _height, _framerate) do
     provided_bitrate
   end
 

diff --git a/lib/membrane_vpx/encoder/vpx_encoder_native.ex b/lib/membrane_vpx/encoder/vpx_encoder_native.ex
@@ -2,9 +2,16 @@ defmodule Membrane.VPx.Encoder.Native do
   @moduledoc false
   use Unifex.Loader
 
-  @spec create!(:vp8 | :vp9, Membrane.VPx.Encoder.encoder_options()) :: reference()
-  def create!(codec, encoder_options) do
-    case create(codec, encoder_options) do
+  @spec create!(
+          :vp8 | :vp9,
+          pos_integer(),
+          pos_integer(),
+          Membrane.RawVideo.pixel_format(),
+          non_neg_integer(),
+          Membrane.VPx.Encoder.user_encoder_config()
+        ) :: reference()
+  def create!(codec, width, height, pixel_format, encoding_deadline, user_encoder_config) do
+    case create(codec, width, height, pixel_format, encoding_deadline, user_encoder_config) 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/keyframes_test.exs b/test/membrane_vpx_plugin/encoder/keyframes_test.exs
@@ -41,19 +41,21 @@ defmodule Membrane.VPx.KeyframesTest do
     test "VP8 codec" do
       perform_test(
         "ref_vp8.raw",
-        %Membrane.VP8.Encoder{encoding_deadline: 1}
+        %Membrane.VP8.Encoder{encoding_deadline: 1, g_lag_in_frames: 0},
+        :vp8
       )
     end
 
     test "VP9 codec" do
       perform_test(
         "ref_vp9.raw",
-        %Membrane.VP9.Encoder{encoding_deadline: 1}
+        %Membrane.VP9.Encoder{encoding_deadline: 1, g_lag_in_frames: 0},
+        :vp9
       )
     end
   end
 
-  defp perform_test(input_file, encoder_struct) do
+  defp perform_test(input_file, encoder_struct, metadata_key) do
     pid =
       Membrane.Testing.Pipeline.start_link_supervised!(
         spec:
@@ -64,14 +66,20 @@ defmodule Membrane.VPx.KeyframesTest do
             pixel_format: :I420,
             width: 1080,
             height: 720,
-            framerate: {30, 1}
+            framerate: {5, 1}
           })
           |> child(:realtimer, Membrane.Realtimer)
           |> child(:encoder, encoder_struct)
           |> child(:keyframe_forcer, KeyframeRequester)
           |> child(:sink, Membrane.Testing.Sink)
       )
 
+    Enum.each(1..5, fn _n ->
+      assert_sink_buffer(pid, :sink, %Membrane.Buffer{
+        metadata: %{^metadata_key => %{is_keyframe: true}}
+      })
+    end)
+
     assert_end_of_stream(pid, :sink, :input, 10_000)
 
     Membrane.Testing.Pipeline.terminate(pid)