Add output_alignment option

lib/membrane_h264_plugin/parser.ex

@@ -55,7 +55,8 @@ defmodule Membrane.H264.Parser do
 
   def_output_pad :output,
     demand_mode: :auto,
-    accepted_format: %H264{alignment: :au, nalu_in_metadata?: true}
+    accepted_format:
+      any_of(%H264{alignment: :au, nalu_in_metadata?: true}, %H264{alignment: :nalu})
 
   def_options sps: [
                 spec: binary(),
@@ -82,6 +83,16 @@ defmodule Membrane.H264.Parser do
                 Its value will be sent inside the output Membrane.H264 stream format.
                 """
               ],
+              output_alignment: [
+                spec: :au | :nalu,
+                default: :au,
+                description: """
+                Alignment of the buffers produced as an output of the parser.
+                If set to `:au`, each output buffer will be a single access unit.
+                Otherwise, if set to `:nalu`, each output buffer will be a single NAL unit.
+                Defaults to `:au`.
+                """
+              ],
               repeat_parameter_sets: [
                 spec: boolean(),
                 default: false,
@@ -107,6 +118,7 @@ defmodule Membrane.H264.Parser do
       previous_timestamps: {nil, nil},
       framerate: opts.framerate,
       au_counter: 0,
+      output_alignment: opts.output_alignment,
       frame_prefix: <<>>,
       parameter_sets_present?: byte_size(opts.sps) > 0 or byte_size(opts.pps) > 0,
       repeat_parameter_sets?: opts.repeat_parameter_sets,
@@ -262,8 +274,10 @@ defmodule Membrane.H264.Parser do
 
         {pts, dts} = prepare_timestamps(buffer_pts, buffer_dts, state, profile, cnt)
 
-        {actions_acc ++ sps_actions ++ [{:buffer, {:output, wrap_into_buffer(au, pts, dts)}}],
-         state, cnt + 1, profile}
+        {actions_acc ++
+           sps_actions ++
+           [{:buffer, {:output, wrap_into_buffer(au, pts, dts, state.output_alignment)}}], state,
+         cnt + 1, profile}
       end)
 
     state = %{state | profile: profile, au_counter: au_counter}
@@ -284,7 +298,12 @@ defmodule Membrane.H264.Parser do
         {[], profile}
 
       sps_nalu ->
-        fmt = Format.from_sps(sps_nalu, framerate: state.framerate)
+        fmt =
+          Format.from_sps(sps_nalu,
+            framerate: state.framerate,
+            output_alignment: state.output_alignment
+          )
+
         {[stream_format: {:output, fmt}], fmt.profile}
     end
   end
@@ -353,8 +372,8 @@ defmodule Membrane.H264.Parser do
 
   defp idr_au?(au), do: :idr in Enum.map(au, & &1.type)
 
-  defp wrap_into_buffer(access_unit, pts, dts) do
-    metadata = prepare_metadata(access_unit)
+  defp wrap_into_buffer(access_unit, pts, dts, :au) do
+    metadata = prepare_au_metadata(access_unit)
 
     buffer =
       access_unit
@@ -368,8 +387,16 @@ defmodule Membrane.H264.Parser do
     buffer
   end
 
-  defp prepare_metadata(nalus) do
-    is_keyframe = Enum.any?(nalus, fn nalu -> nalu.type == :idr end)
+  defp wrap_into_buffer(access_unit, pts, dts, :nalu) do
+    access_unit
+    |> Enum.zip(prepare_nalus_metadata(access_unit))
+    |> Enum.map(fn {nalu, metadata} ->
+      %Buffer{payload: nalu.payload, metadata: metadata, pts: pts, dts: dts}
+    end)
+  end
+
+  defp prepare_au_metadata(nalus) do
+    is_keyframe? = Enum.any?(nalus, fn nalu -> nalu.type == :idr end)
 
     nalus =
       nalus
@@ -395,7 +422,7 @@ defmodule Membrane.H264.Parser do
 
         metadata =
           if i == 0 do
-            put_in(metadata, [:metadata, :h264, :new_access_unit], %{key_frame?: is_keyframe})
+            put_in(metadata, [:metadata, :h264, :new_access_unit], %{key_frame?: is_keyframe?})
           else
             metadata
           end
@@ -404,7 +431,21 @@ defmodule Membrane.H264.Parser do
       end)
       |> elem(0)
 
-    %{h264: %{key_frame?: is_keyframe, nalus: nalus}}
+    %{h264: %{key_frame?: is_keyframe?, nalus: nalus}}
+  end
+
+  defp prepare_nalus_metadata(nalus) do
+    is_keyframe? = Enum.any?(nalus, fn nalu -> nalu.type == :idr end)
+
+    Enum.with_index(nalus)
+    |> Enum.map(fn {nalu, i} ->
+      %{h264: %{type: nalu.type}}
+      |> Bunch.then_if(
+        i == 0,
+        &put_in(&1, [:h264, :new_access_unit], %{key_frame?: is_keyframe?})
+      )
+      |> Bunch.then_if(i == length(nalus) - 1, &put_in(&1, [:h264, :end_access_unit], true))
+    end)
   end
 
   defp stream_format_sent?(actions, %{pads: %{output: %{stream_format: nil}}}),

lib/membrane_h264_plugin/parser/format.ex

@@ -6,15 +6,6 @@ defmodule Membrane.H264.Parser.Format do
 
   alias Membrane.H264
 
-  @default_format %H264{
-    alignment: :au,
-    framerate: nil,
-    height: 720,
-    nalu_in_metadata?: true,
-    profile: :high,
-    width: 1280
-  }
-
   @profiles_description [
     high_cavlc_4_4_4_intra: [profile_idc: 44],
     constrained_baseline: [profile_idc: 66, constraint_set1: 1],
@@ -39,7 +30,10 @@ defmodule Membrane.H264.Parser.Format do
   """
   @spec from_sps(
           sps_nalu :: H264.Parser.NALu.t(),
-          options_fields :: [framerate: {pos_integer(), pos_integer()}]
+          options_fields :: [
+            framerate: {pos_integer(), pos_integer()},
+            output_alignment: :au | :nalu
+          ]
         ) :: H264.t()
   def from_sps(sps_nalu, options_fields) do
     sps = sps_nalu.parsed_fields
@@ -77,11 +71,12 @@ defmodule Membrane.H264.Parser.Format do
     profile = parse_profile(sps_nalu)
 
     %H264{
-      @default_format
-      | width: width,
-        height: height,
-        profile: profile,
-        framerate: Keyword.get(options_fields, :framerate)
+      width: width,
+      height: height,
+      profile: profile,
+      framerate: Keyword.get(options_fields, :framerate),
+      alignment: Keyword.get(options_fields, :output_alignment),
+      nalu_in_metadata?: true
     }
   end
 

test/integration/modes_test.exs

@@ -99,4 +99,30 @@ defmodule Membrane.H264.ModesTest do
 
     Pipeline.terminate(pid, blocking?: true)
   end
+
+  test "if single NAL unit is sent per buffer with `output_alignment: :nalu`" do
+    {:ok, _supervisor_pid, pid} =
+      Pipeline.start_supervised(
+        structure: [
+          child(:source, %Membrane.File.Source{location: @h264_input_file})
+          |> child(:parser, %Parser{output_alignment: :nalu})
+          |> child(:sink, Sink)
+        ]
+      )
+
+    assert_pipeline_play(pid)
+    assert_sink_stream_format(pid, :sink, %Membrane.H264{alignment: :nalu})
+
+    binary = File.read!(@h264_input_file)
+    ref_buffers = prepare_buffers(binary, :nalu_aligned)
+
+    Enum.each(ref_buffers, fn ref_buffer ->
+      assert_sink_buffer(pid, :sink, buffer)
+      assert buffer.payload == ref_buffer.payload
+      assert Map.has_key?(buffer.metadata, :h264) and Map.has_key?(buffer.metadata.h264, :type)
+    end)
+
+    assert_end_of_stream(pid, :sink)
+    Pipeline.terminate(pid, blocking?: true)
+  end
 end