Allow requesting keyframes and add keyframe information in metadata (#4)

* Add is_keyframe field to metadata

* Add docs for metadata

* Remove forcing keyframe from flush

.circleci/config.yml

@@ -10,12 +10,15 @@ workflows:
           filters: &filters
             tags:
               only: /v.*/
+          cache-version: 2
       - elixir/test:
           filters:
             <<: *filters
+          cache-version: 2
       - elixir/lint:
           filters:
             <<: *filters
+          cache-version: 2
       - elixir/hex_publish:
           requires:
             - elixir/build_test
@@ -28,3 +31,4 @@ workflows:
               ignore: /.*/
             tags:
               only: /v.*/
+          cache-version: 2

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,34 @@ 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, State *state) { return encode(env, NULL, 0, 0, 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]}
+       {: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.KeyframeRequestEvent` 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.KeyframeRequestEvent` 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