Allow requesting keyframes and add keyframe information in metadata

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)
@@ -78,30 +81,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);
 }
 
-UNIFEX_TERM encode(UnifexEnv *env, vpx_image_t *img, vpx_codec_pts_t pts, State *state) {
+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, 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 +130,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 +139,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,6 +6,12 @@ type codec :: :vp8 | :vp9
 
 type pixel_format :: :I420 | :I422 | :I444 | :NV12 | :YV12
 
+type encoded_frame :: %EncodedFrame{
+       payload: payload,
+       pts: int64,
+       is_keyframe: bool
+     }
+
 spec create(
        codec,
        width :: unsigned,
@@ -15,12 +21,12 @@ spec create(
      ) ::
        {: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
 
@@ -36,6 +45,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
 
@@ -36,6 +45,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

lib/membrane_vpx/encoder/vpx_encoder.ex

@@ -1,7 +1,7 @@
 defmodule Membrane.VPx.Encoder do
   @moduledoc false
 
-  alias Membrane.{Buffer, RawVideo, VP8, VP9}
+  alias Membrane.{Buffer, KeyframeRequestEvent, RawVideo, VP8, VP9}
   alias Membrane.Element.CallbackContext
   alias Membrane.VPx.Encoder.Native
 
@@ -14,20 +14,23 @@ defmodule Membrane.VPx.Encoder do
             codec: :vp8 | :vp9,
             codec_module: VP8 | VP9,
             encoding_deadline: non_neg_integer(),
-            encoder_ref: reference() | nil
+            encoder_ref: reference() | nil,
+            force_next_keyframe: boolean()
           }
 
     @enforce_keys [:codec, :codec_module, :encoding_deadline]
     defstruct @enforce_keys ++
                 [
-                  encoder_ref: nil
+                  encoder_ref: nil,
+                  force_next_keyframe: false
                 ]
   end
 
   @type callback_return :: {[Membrane.Element.Action.t()], State.t()}
 
   @spec handle_init(CallbackContext.t(), VP8.Encoder.t() | VP9.Encoder.t(), :vp8 | :vp9) ::
           callback_return()
+
   def handle_init(_ctx, opts, codec) do
     state = %State{
       codec: codec,
@@ -57,27 +60,34 @@ defmodule Membrane.VPx.Encoder do
     {flushed_buffers, encoder_ref} =
       maybe_recreate_encoder(ctx.pads.input.stream_format, stream_format, state)
 
+    force_next_keyframe = if flushed_buffers == [], do: state.force_next_keyframe, else: false
+
     {
       [buffer: {:output, flushed_buffers}, stream_format: {:output, output_stream_format}],
-      %{state | encoder_ref: encoder_ref}
+      %{state | encoder_ref: encoder_ref, force_next_keyframe: force_next_keyframe}
     }
   end
 
   @spec handle_buffer(:input, Membrane.Buffer.t(), CallbackContext.t(), State.t()) ::
           callback_return()
   def handle_buffer(:input, %Buffer{payload: payload, pts: pts}, _ctx, state) do
-    {:ok, encoded_frames, timestamps} = Native.encode_frame(payload, pts, state.encoder_ref)
+    {:ok, encoded_frames} =
+      Native.encode_frame(payload, pts, state.force_next_keyframe, state.encoder_ref)
 
-    buffers =
-      Enum.zip(encoded_frames, timestamps)
-      |> Enum.map(fn {frame, frame_pts} -> %Buffer{payload: frame, pts: frame_pts} end)
+    buffers = get_buffers_from_frames(encoded_frames, state.codec)
+
+    {[buffer: {:output, buffers}], %{state | force_next_keyframe: false}}
+  end
 
-    {[buffer: {:output, buffers}], state}
+  @spec handle_event(:output, KeyframeRequestEvent.t(), CallbackContext.t(), State.t()) ::
+          callback_return()
+  def handle_event(:output, %KeyframeRequestEvent{}, _ctx, state) do
+    {[], %{state | force_next_keyframe: true}}
   end
 
   @spec handle_end_of_stream(:input, CallbackContext.t(), State.t()) :: callback_return()
   def handle_end_of_stream(:input, _ctx, state) do
-    buffers = flush(state.encoder_ref)
+    buffers = flush(state.force_next_keyframe, state.encoder_ref, state.codec)
     {[buffer: {:output, buffers}, end_of_stream: :output], state}
   end
 
@@ -109,16 +119,29 @@ defmodule Membrane.VPx.Encoder do
       Native.create!(state.codec, width, height, pixel_format, encoding_deadline)
 
     case state.encoder_ref do
-      nil -> {[], new_encoder_ref}
-      old_encoder_ref -> {flush(old_encoder_ref), new_encoder_ref}
+      nil ->
+        {[], new_encoder_ref}
+
+      old_encoder_ref ->
+        {flush(state.force_next_keyframe, old_encoder_ref, state.codec), new_encoder_ref}
     end
   end
 
-  @spec flush(reference()) :: [Membrane.Buffer.t()]
-  defp flush(encoder_ref) do
-    {:ok, encoded_frames, timestamps} = Native.flush(encoder_ref)
+  @spec flush(boolean(), reference(), :vp8 | :vp9) :: [Membrane.Buffer.t()]
+  defp flush(force_next_keyframe, encoder_ref, codec) do
+    {:ok, encoded_frames} = Native.flush(force_next_keyframe, encoder_ref)
+
+    get_buffers_from_frames(encoded_frames, codec)
+  end
 
-    Enum.zip(encoded_frames, timestamps)
-    |> Enum.map(fn {frame, frame_pts} -> %Buffer{payload: frame, pts: frame_pts} end)
+  @spec get_buffers_from_frames([EncodedFrame.t()], :vp8 | :vp9) :: [Buffer.t()]
+  def get_buffers_from_frames(encoded_frames, codec) do
+    Enum.map(encoded_frames, fn %{payload: payload, pts: pts, is_keyframe: is_keyframe} ->
+      %Buffer{
+        payload: payload,
+        pts: pts,
+        metadata: %{codec => %{is_keyframe: is_keyframe}}
+      }
+    end)
   end
 end