Decoder prototype working

c_src/membrane_vpx_plugin/vp8_decoder.c

@@ -1,17 +1,114 @@
 #include "vp8_decoder.h"
 
+void handle_destroy_state(UnifexEnv *env, State *state) {
+  UNIFEX_UNUSED(env);
+
+  vpx_codec_destroy(&state->codec_context);
+}
+
 UNIFEX_TERM create(UnifexEnv *env) {
   UNIFEX_TERM result;
   State *state = unifex_alloc_state(env);
-  vpx_codec_ctx_t codec_context;
-  vpx_codec_iface_t *codec_interface;
+  state->codec_interface = vpx_codec_vp8_dx();
 
-  if (vpx_codec_dec_init(&codec_context, codec_interface, NULL, 0)) {
+  if (vpx_codec_dec_init(&state->codec_context, state->codec_interface, NULL,
+                         0)) {
     result = create_result_error(env, "Failed to initialize decoder");
     unifex_release_state(env, state);
     return result;
   }
   result = create_result_ok(env, state);
   unifex_release_state(env, state);
   return result;
-}
+}
+
+Dimensions get_plane_dimensions(const vpx_image_t *img, int plane) {
+  const int height = (plane > 0 && img->y_chroma_shift > 0)
+                         ? (img->d_h + 1) >> img->y_chroma_shift
+                         : img->d_h;
+
+  int width = (plane > 0 && img->x_chroma_shift > 0)
+                  ? (img->d_w + 1) >> img->x_chroma_shift
+                  : img->d_w;
+
+  // Fixing NV12 chroma width if it is odd
+  if (img->fmt == VPX_IMG_FMT_NV12 && plane == 1)
+    width = (width + 1) & ~1;
+
+  return (Dimensions){width, height};
+}
+size_t get_image_byte_size(const vpx_image_t *img) {
+  const int bytes_per_pixel = (img->fmt & VPX_IMG_FMT_HIGHBITDEPTH) ? 2 : 1;
+  const int number_of_planes = (img->fmt == VPX_IMG_FMT_NV12) ? 2 : 3;
+
+  size_t image_size = 0;
+
+  for (int plane = 0; plane < number_of_planes; ++plane) {
+    Dimensions plane_dimensions = get_plane_dimensions(img, plane);
+    image_size +=
+        plane_dimensions.width * plane_dimensions.height * bytes_per_pixel;
+  }
+  return image_size;
+}
+
+void get_raw_frame_from_image(const vpx_image_t *img,
+                              UnifexPayload *output_frame) {
+  const int bytes_per_pixel = (img->fmt & VPX_IMG_FMT_HIGHBITDEPTH) ? 2 : 1;
+
+  // Assuming that for nv12 we write all chroma data at once
+  const int number_of_planes = (img->fmt == VPX_IMG_FMT_NV12) ? 2 : 3;
+  unsigned char *output_data = output_frame->data;
+
+  for (int plane = 0; plane < number_of_planes; ++plane) {
+    const unsigned char *buf = img->planes[plane];
+    const int stride = img->stride[plane];
+    Dimensions plane_dimensions = get_plane_dimensions(img, plane);
+
+    for (unsigned int y = 0; y < plane_dimensions.height; ++y) {
+      size_t bytes_to_write = bytes_per_pixel * plane_dimensions.width;
+      memcpy(output_data, buf, bytes_to_write);
+      buf += stride;
+      output_data += bytes_to_write;
+    }
+  }
+}
+
+UNIFEX_TERM decode_frame(UnifexEnv *env, UnifexPayload *frame, State *state) {
+  vpx_codec_iter_t iter = NULL;
+  vpx_image_t *img = NULL;
+  unsigned int frames_cnt = 0, max_frames = 2;
+  UnifexPayload **output_frames =
+      unifex_alloc(max_frames * sizeof(*output_frames));
+
+  if (vpx_codec_decode(&state->codec_context, frame->data, frame->size, NULL,
+                       0)) {
+    return decode_frame_result_error(env, "Decoding frame failed");
+  }
+
+  while ((img = vpx_codec_get_frame(&state->codec_context, &iter)) != NULL) {
+    if (frames_cnt >= max_frames) {
+      max_frames *= 2;
+      output_frames =
+          unifex_realloc(output_frames, max_frames * sizeof(*output_frames));
+    }
+
+    output_frames[frames_cnt] = unifex_alloc(sizeof(UnifexPayload));
+    const size_t output_frame_size = get_image_byte_size(img);
+    unifex_payload_alloc(env, UNIFEX_PAYLOAD_BINARY, output_frame_size,
+                         output_frames[frames_cnt]);
+
+    get_raw_frame_from_image(img, output_frames[frames_cnt]);
+    frames_cnt++;
+  }
+
+  UNIFEX_TERM result = decode_frame_result_ok(env, output_frames, frames_cnt);
+  for (unsigned int i = 0; i < frames_cnt; i++) {
+    if (output_frames[i] != NULL) {
+      unifex_payload_release(output_frames[i]);
+      unifex_free(output_frames[i]);
+    }
+  }
+  unifex_free(output_frames);
+
+  return result;
+}

c_src/membrane_vpx_plugin/vp8_decoder.h

@@ -1,8 +1,16 @@
-#include "_generated/vp8_decoder.h"
+#pragma once
 #include "vpx/vp8dx.h"
 #include "vpx/vpx_decoder.h"
+#include <erl_nif.h>
 
 typedef struct State {
   vpx_codec_ctx_t codec_context;
   vpx_codec_iface_t *codec_interface;
-} State;
+} State;
+
+typedef struct Dimensions {
+  unsigned int width;
+  unsigned int height;
+} Dimensions;
+
+#include "_generated/vp8_decoder.h"

c_src/membrane_vpx_plugin/vp8_decoder.spec.exs

@@ -6,3 +6,5 @@ spec create() :: {:ok :: label, state} | {:error :: label, reason :: atom}
 
 spec decode_frame(payload, state) ::
        {:ok :: label, frames :: [payload]} | {:error :: label, reason :: atom}
+
+dirty :cpu, decode_frame: 2

lib/membrane_vpx/vp8_decoder.ex

@@ -4,27 +4,68 @@ defmodule Membrane.VP8.Decoder do
   """
   use Membrane.Filter
 
+  alias Membrane.Buffer
+  alias Membrane.VP8
+  alias __MODULE__.Native
+
   def_input_pad :input,
-    accepted_format: Membrane.VP8
+    accepted_format:
+      any_of(
+        Membrane.VP8,
+        %Membrane.RemoteStream{content_format: format} when format in [nil, VP8]
+      )
 
   def_output_pad :output,
     accepted_format: Membrane.RawVideo
 
   defmodule State do
     @moduledoc false
     @type t :: %__MODULE__{
-            native: term()
+            decoder_ref: reference() | nil
           }
 
     @enforce_keys []
     defstruct @enforce_keys ++
                 [
-                  native: nil
+                  decoder_ref: nil
                 ]
   end
 
   @impl true
   def handle_init(_ctx, _opts) do
     {[], %State{}}
   end
+
+  @impl true
+  def handle_setup(_ctx, state) do
+    native = Native.create!()
+
+    {[], %{state | decoder_ref: native}}
+  end
+
+  @impl true
+  def handle_playing(_ctx, state) do
+    stream_format =
+      %Membrane.RawVideo{
+        width: 0,
+        height: 0,
+        framerate: {0, 0},
+        pixel_format: :I420,
+        aligned: true
+      }
+
+    {[stream_format: {:output, stream_format}], state}
+  end
+
+  @impl true
+  def handle_stream_format(:input, _stream_format, _ctx, state) do
+    {[], state}
+  end
+
+  @impl true
+  def handle_buffer(:input, %Buffer{payload: payload}, _ctx, state) do
+    {:ok, decoded_frames} = Native.decode_frame(payload, state.decoder_ref)
+    buffers = Enum.map(decoded_frames, &%Buffer{payload: &1})
+    {[buffer: {:output, buffers}], state}
+  end
 end

lib/membrane_vpx/vp8_decoder_native.ex

@@ -1,4 +1,12 @@
 defmodule Membrane.VP8.Decoder.Native do
   @moduledoc false
   use Unifex.Loader
+
+  @spec create! :: reference()
+  def create!() do
+    case create() do
+      {:ok, decoder_ref} -> decoder_ref
+      {:error, reason} -> raise "Failed to create native decoder: #{inspect(reason)}"
+    end
+  end
 end

mix.exs

@@ -41,6 +41,9 @@ defmodule Membrane.VPx.Plugin.Mixfile do
       {:membrane_core, "~> 1.0"},
       {:unifex, "~> 1.2"},
       {:membrane_raw_video_format, "~> 0.4.0"},
+      {:membrane_vp8_format, "~> 0.4.0"},
+      {:membrane_ivf_plugin, "~> 0.7.0", only: :test},
+      {:membrane_file_plugin, "~> 0.17.0", only: :test},
       {:ex_doc, ">= 0.0.0", only: :dev, runtime: false},
       {:dialyxir, ">= 0.0.0", only: :dev, runtime: false},
       {:credo, ">= 0.0.0", only: :dev, runtime: false}