Add functionality of converting between avc1, avc3 and Annex B stream…

… structures (#36)

* Add functionality of converting between stream structures

* Add stream structure conversion tests

* Make NALu structs hold unprefixed NALu payloads

* Bump version to 0.7.0

README.md

@@ -22,7 +22,7 @@ The package can be installed by adding `membrane_h264_plugin` to your list of de
 ```elixir
 def deps do
   [
-	  {:membrane_h264_plugin, "~> 0.6.0"}
+	  {:membrane_h264_plugin, "~> 0.7.0"}
   ]
 end
 ```

lib/membrane_h264_plugin/parser/au_splitter.ex

@@ -28,7 +28,7 @@ defmodule Membrane.H264.Parser.AUSplitter do
             nalus_acc: [NALu.t()],
             fsm_state: :first | :second,
             previous_primary_coded_picture_nalu: NALu.t() | nil,
-            access_units_to_output: access_unit_t()
+            access_units_to_output: access_unit()
           }
   @enforce_keys [
     :nalus_acc,
@@ -58,7 +58,7 @@ defmodule Membrane.H264.Parser.AUSplitter do
   @typedoc """
   A type representing an access unit - a list of logically associated NAL units.
   """
-  @type access_unit_t() :: list(NALu.t())
+  @type access_unit() :: list(NALu.t())
 
   @doc """
   Splits the given list of NAL units into the access units.
@@ -77,7 +77,7 @@ defmodule Membrane.H264.Parser.AUSplitter do
   is not returned until another access unit starts, as it's the only way to prove that
   the access unit is complete.
   """
-  @spec split([NALu.t()], assume_au_aligned :: boolean(), t()) :: {[access_unit_t()], t()}
+  @spec split([NALu.t()], boolean(), t()) :: {[access_unit()], t()}
   def split(nalus, assume_au_aligned \\ false, state) do
     state = do_split(nalus, state)
 
@@ -112,7 +112,10 @@ defmodule Membrane.H264.Parser.AUSplitter do
         )
 
       true ->
-        Membrane.Logger.warning("AUSplitter: Improper transition")
+        Membrane.Logger.warning(
+          "AUSplitter: Improper transition, first_nalu: #{inspect(first_nalu)}"
+        )
+
         state
     end
   end
@@ -157,7 +160,10 @@ defmodule Membrane.H264.Parser.AUSplitter do
         )
 
       true ->
-        Membrane.Logger.warning("AUSplitter: Improper transition")
+        Membrane.Logger.warning(
+          "AUSplitter: Improper transition, first_nalu: #{inspect(first_nalu)}"
+        )
+
         state
     end
   end

lib/membrane_h264_plugin/parser/decoder_configuration_record.ex

@@ -1,58 +1,94 @@
 defmodule Membrane.H264.Parser.DecoderConfigurationRecord do
   @moduledoc """
-  Utility functions for parsing AVC Configuration Record.
+  Utility functions for parsing and generating AVC Configuration Record.
 
   The structure of the record is described in section 5.2.4.1.1 of MPEG-4 part 15 (ISO/IEC 14496-15).
   """
+
+  alias Membrane.H264.Parser
+
   @enforce_keys [
-    :sps,
-    :pps,
+    :spss,
+    :ppss,
     :avc_profile_indication,
     :avc_level,
     :profile_compatibility,
-    :length_size_minus_one
+    :nalu_length_size
   ]
   defstruct @enforce_keys
 
   @typedoc "Structure representing the Decoder Configuartion Record"
   @type t() :: %__MODULE__{
-          sps: [binary()],
-          pps: [binary()],
+          spss: [binary()],
+          ppss: [binary()],
           avc_profile_indication: non_neg_integer(),
           profile_compatibility: non_neg_integer(),
           avc_level: non_neg_integer(),
-          length_size_minus_one: non_neg_integer()
+          nalu_length_size: pos_integer()
         }
 
+  @doc """
+  Generates a DCR based on given PPSs and SPSs.
+  """
+  @spec generate([binary()], [binary()], Parser.stream_structure()) ::
+          binary() | nil
+  def generate([], _ppss, _stream_structure) do
+    nil
+  end
+
+  def generate(spss, ppss, {avc, nalu_length_size}) do
+    <<_idc_and_type, profile, compatibility, level, _rest::binary>> = List.last(spss)
+
+    cond do
+      avc == :avc1 ->
+        <<1, profile, compatibility, level, 0b111111::6, nalu_length_size - 1::2-integer,
+          0b111::3, length(spss)::5-integer, encode_parameter_sets(spss)::binary,
+          length(ppss)::8-integer, encode_parameter_sets(ppss)::binary>>
+
+      avc == :avc3 ->
+        <<1, profile, compatibility, level, 0b111111::6, nalu_length_size - 1::2-integer,
+          0b111::3, 0::5, 0::8>>
+    end
+  end
+
+  defp encode_parameter_sets(pss) do
+    Enum.map_join(pss, &<<byte_size(&1)::16-integer, &1::binary>>)
+  end
+
+  @spec remove_parameter_sets(binary()) :: binary()
+  def remove_parameter_sets(dcr) do
+    <<dcr_head::binary-(8 * 5), _rest::binary>> = dcr
+    <<dcr_head::binary, 0b111::3, 0::5, 0::8>>
+  end
+
   @doc """
   Parses the DCR.
   """
-  @spec parse(binary()) :: {:ok, t()} | {:error, any()}
+  @spec parse(binary()) :: t()
   def parse(
         <<1::8, avc_profile_indication::8, profile_compatibility::8, avc_level::8, 0b111111::6,
           length_size_minus_one::2, 0b111::3, rest::bitstring>>
       ) do
-    {sps, rest} = parse_sps(rest)
-    {pps, _rest} = parse_pps(rest)
+    {spss, rest} = parse_spss(rest)
+    {ppss, _rest} = parse_ppss(rest)
 
     %__MODULE__{
-      sps: sps,
-      pps: pps,
+      spss: spss,
+      ppss: ppss,
       avc_profile_indication: avc_profile_indication,
       profile_compatibility: profile_compatibility,
       avc_level: avc_level,
-      length_size_minus_one: length_size_minus_one
+      nalu_length_size: length_size_minus_one + 1
     }
-    |> then(&{:ok, &1})
   end
 
   def parse(_data), do: {:error, :unknown_pattern}
 
-  defp parse_sps(<<num_of_sps::5, rest::bitstring>>) do
-    do_parse_array(num_of_sps, rest)
+  defp parse_spss(<<num_of_spss::5, rest::bitstring>>) do
+    do_parse_array(num_of_spss, rest)
   end
 
-  defp parse_pps(<<num_of_pps::8, rest::bitstring>>), do: do_parse_array(num_of_pps, rest)
+  defp parse_ppss(<<num_of_ppss::8, rest::bitstring>>), do: do_parse_array(num_of_ppss, rest)
 
   defp do_parse_array(amount, rest, acc \\ [])
   defp do_parse_array(0, rest, acc), do: {Enum.reverse(acc), rest}

lib/membrane_h264_plugin/parser/format.ex

@@ -30,9 +30,10 @@ defmodule Membrane.H264.Parser.Format do
   """
   @spec from_sps(
           sps_nalu :: H264.Parser.NALu.t(),
+          output_raw_stream_structure :: H264.stream_structure(),
           options_fields :: [output_alignment: :au | :nalu]
         ) :: H264.t()
-  def from_sps(sps_nalu, options_fields) do
+  def from_sps(sps_nalu, output_raw_stream_structure, options_fields) do
     sps = sps_nalu.parsed_fields
 
     chroma_array_type = if sps.separate_colour_plane_flag == 0, do: sps.chroma_format_idc, else: 0
@@ -72,7 +73,8 @@ defmodule Membrane.H264.Parser.Format do
       height: height,
       profile: profile,
       alignment: Keyword.get(options_fields, :output_alignment),
-      nalu_in_metadata?: true
+      nalu_in_metadata?: true,
+      stream_structure: output_raw_stream_structure
     }
   end
 

lib/membrane_h264_plugin/parser/nalu.ex

@@ -10,24 +10,24 @@ defmodule Membrane.H264.Parser.NALu do
   In the structure there ardqde following fields:
   * `parsed_fields` - the map with keys being the NALu field names and the values being the value fetched from the NALu binary.
   They correspond to the NALu schemes defined in the section *7.3.* of the *"ITU-T Rec. H.264 (01/2012)"*.
-  * `prefix_length` - number of bytes of the prefix used to split the NAL units in the bytestream.
-  The prefix is defined as in: *"Annex B"* of the *"ITU-T Rec. H.264 (01/2012)"*.
+  * `stripped_prefix` - prefix that used to split the NAL units in the bytestream and was stripped from the payload.
+  The prefix is defined as in: *"Annex B"* of the *"ISO/IEC 14496-10"* or in "ISO/IEC 14496-15".
   * `type` - an atom representing the type of the NALu. Atom's name is based on the
   *"Table 7-1 – NAL unit type codes, syntax element categories, and NAL unit type classes"* of the *"ITU-T Rec. H.264 (01/2012)"*.
-  * `payload` - the binary, which parsing resulted in that structure being produced
+  * `payload` - the binary, which parsing resulted in that structure being produced stripped of it's prefix
   * `status` - `:valid`, if the parsing was successfull, `:error` otherwise
   """
   @type t :: %__MODULE__{
           parsed_fields: %{atom() => any()},
-          prefix_length: pos_integer(),
           type: Membrane.H264.Parser.NALuTypes.nalu_type(),
+          stripped_prefix: binary(),
           payload: binary(),
           status: :valid | :error,
           timestamps: timestamps()
         }
 
   @type timestamps :: {pts :: integer() | nil, dts :: integer() | nil}
 
-  @enforce_keys [:parsed_fields, :prefix_length, :type, :payload, :status]
+  @enforce_keys [:parsed_fields, :type, :stripped_prefix, :payload, :status]
   defstruct @enforce_keys ++ [timestamps: {nil, nil}]
 end

lib/membrane_h264_plugin/parser/nalu_parser.ex

@@ -3,55 +3,63 @@ defmodule Membrane.H264.Parser.NALuParser do
   A module providing functionality of parsing a stream of binaries, out of which each
   is a payload of a single NAL unit.
   """
+
+  alias Membrane.H264.Parser
   alias Membrane.H264.Parser.{NALu, NALuTypes}
-  alias Membrane.H264.Parser.NALuParser.SchemeParser
-  alias Membrane.H264.Parser.NALuParser.Schemes
+  alias Membrane.H264.Parser.NALuParser.{SchemeParser, Schemes}
+
+  @annexb_prefix_code <<0, 0, 0, 1>>
 
   @typedoc """
   A structure holding the state of the NALu parser.
   """
   @opaque t :: %__MODULE__{
-            scheme_parser_state: SchemeParser.t()
+            scheme_parser_state: SchemeParser.t(),
+            input_stream_structure: Parser.stream_structure()
           }
-
-  defstruct scheme_parser_state: SchemeParser.new()
+  @enforce_keys [:input_stream_structure]
+  defstruct @enforce_keys ++ [scheme_parser_state: SchemeParser.new()]
 
   @doc """
-  Returns a structure holding a clear NALu parser state.
+  Returns a structure holding a clear NALu parser state. `input_stream_structure`
+  determines the prefixes of input NALU payloads.
   """
-  @spec new() :: t()
-  def new(), do: %__MODULE__{}
+  @spec new(Parser.stream_structure()) :: t()
+  def new(input_stream_structure \\ :annexb) do
+    %__MODULE__{input_stream_structure: input_stream_structure}
+  end
 
   @doc """
   Parses a list of binaries, each representing a single NALu.
 
   See `parse/3` for details.
   """
-  @spec parse_nalus([binary()], NALu.timestamps(), t()) :: {[NALu.t()], t()}
-  def parse_nalus(nalus_payloads, timestamps \\ {nil, nil}, state) do
+  @spec parse_nalus([binary()], NALu.timestamps(), boolean(), t()) :: {[NALu.t()], t()}
+  def parse_nalus(nalus_payloads, timestamps \\ {nil, nil}, payload_prefixed? \\ true, state) do
     Enum.map_reduce(nalus_payloads, state, fn nalu_payload, state ->
-      parse(nalu_payload, timestamps, state)
+      parse(nalu_payload, timestamps, payload_prefixed?, state)
     end)
   end
 
   @doc """
-  Parses a binary representing a single NALu.
+  Parses a binary representing a single NALu and removes it's prefix (if it exists).
 
   Returns a structure that
   contains parsed fields fetched from that NALu.
-  The input binary is expected to contain the prefix, defined as in
-  the *"Annex B"* of the *"ITU-T Rec. H.264 (01/2012)"*.
+  When `payload_prefixed?` is true the input binary is expected to contain one of:
+  * prefix defined as the *"Annex B"* of the *"ITU-T Rec. H.264 (01/2012)"*.
+  * prefix of size defined in state describing the length of the NALU in bytes, as described in *ISO/IEC 14496-15*.
   """
-  @spec parse(binary(), timestamps :: {pts :: integer() | nil, dts :: integer() | nil}, t()) ::
-          {NALu.t(), t()}
-  def parse(nalu_payload, timestamps \\ {nil, nil}, state) do
-    {prefix_length, nalu_payload_without_prefix} =
-      case nalu_payload do
-        <<0, 0, 1, rest::binary>> -> {3, rest}
-        <<0, 0, 0, 1, rest::binary>> -> {4, rest}
+  @spec parse(binary(), NALu.timestamps(), boolean(), t()) :: {NALu.t(), t()}
+  def parse(nalu_payload, timestamps \\ {nil, nil}, payload_prefixed? \\ true, state) do
+    {prefix, unprefixed_nalu_payload} =
+      if payload_prefixed? do
+        unprefix_nalu_payload(nalu_payload, state.input_stream_structure)
+      else
+        {0, nalu_payload}
       end
 
-    <<nalu_header::binary-size(1), nalu_body::binary>> = nalu_payload_without_prefix
+    <<nalu_header::binary-size(1), nalu_body::binary>> = unprefixed_nalu_payload
 
     new_scheme_parser_state = SchemeParser.new(state.scheme_parser_state)
 
@@ -73,8 +81,8 @@ defmodule Membrane.H264.Parser.NALuParser do
            parsed_fields: parsed_fields,
            type: type,
            status: :valid,
-           prefix_length: prefix_length,
-           payload: nalu_payload,
+           stripped_prefix: prefix,
+           payload: unprefixed_nalu_payload,
            timestamps: timestamps
          }, scheme_parser_state}
       catch
@@ -83,19 +91,65 @@ defmodule Membrane.H264.Parser.NALuParser do
              parsed_fields: parsed_fields,
              type: type,
              status: :error,
-             prefix_length: prefix_length,
-             payload: nalu_payload,
+             stripped_prefix: prefix,
+             payload: unprefixed_nalu_payload,
              timestamps: timestamps
            }, scheme_parser_state}
       end
 
-    state = %__MODULE__{
-      scheme_parser_state: scheme_parser_state
-    }
+    state = %{state | scheme_parser_state: scheme_parser_state}
 
     {nalu, state}
   end
 
+  @doc """
+  Returns payload of the NALu with appropriate prefix generated based on output stream
+  structure and prefix length.
+  """
+  @spec get_prefixed_nalu_payload(NALu.t(), Parser.stream_structure(), boolean()) :: binary()
+  def get_prefixed_nalu_payload(nalu, output_stream_structure, stable_prefixing? \\ true) do
+    case {output_stream_structure, stable_prefixing?} do
+      {:annexb, true} ->
+        case nalu.stripped_prefix do
+          <<0, 0, 1>> -> <<0, 0, 1, nalu.payload::binary>>
+          <<0, 0, 0, 1>> -> <<0, 0, 0, 1, nalu.payload::binary>>
+          _prefix -> @annexb_prefix_code <> nalu.payload
+        end
+
+      {:annexb, false} ->
+        @annexb_prefix_code <> nalu.payload
+
+      {{_avc, nalu_length_size}, _stable_prefixing?} ->
+        <<byte_size(nalu.payload)::integer-size(nalu_length_size)-unit(8), nalu.payload::binary>>
+    end
+  end
+
+  @spec unprefix_nalu_payload(binary(), Parser.stream_structure()) ::
+          {stripped_prefix :: binary(), payload :: binary()}
+  defp unprefix_nalu_payload(nalu_payload, :annexb) do
+    case nalu_payload do
+      <<0, 0, 1, rest::binary>> -> {<<0, 0, 1>>, rest}
+      <<0, 0, 0, 1, rest::binary>> -> {<<0, 0, 0, 1>>, rest}
+    end
+  end
+
+  defp unprefix_nalu_payload(nalu_payload, {_avc, nalu_length_size}) do
+    <<nalu_length::integer-size(nalu_length_size)-unit(8), rest::binary>> = nalu_payload
+
+    {<<nalu_length::integer-size(nalu_length_size)-unit(8)>>, rest}
+  end
+
+  @spec prefix_nalus_payloads([binary()], Parser.stream_structure()) :: binary()
+  def prefix_nalus_payloads(nalus, :annexb) do
+    Enum.join([<<>> | nalus], @annexb_prefix_code)
+  end
+
+  def prefix_nalus_payloads(nalus, {_avc, nalu_length_size}) do
+    Enum.map_join(nalus, fn nalu ->
+      <<byte_size(nalu)::integer-size(nalu_length_size)-unit(8), nalu::binary>>
+    end)
+  end
+
   defp parse_proper_nalu_type(payload, state, type) do
     case type do
       :sps ->