AV1 muxer (#125)

* Support AV1 in Muxer

* added membrane_av1_format library

* removed redundant deps

* fix: reverted unrelated changes, refactored code, use the membrane_av1_format dependency

* fix: addressed the PR comments

* chore: bump membrane_av1_format to 0.3.1

* chore: addressed naming suggestions and fixed credo

* test: added muxer test

* Satisfy credo and compiler

---------

Co-authored-by: noarkhh <jakub.pryc@gmail.com>

Author: maslowalex

lib/membrane_mp4/container/schema_default.ex

@@ -50,6 +50,20 @@ defmodule Membrane.MP4.Container.Schema.Default do
     ]
   ]
 
+  @av1_schema [
+    version: 0,
+    fields: @visual_sample_entry,
+    av1C: [
+      black_box?: true
+    ],
+    pasp: [
+      fields: [
+        h_spacing: :uint32,
+        v_spacing: :uint32
+      ]
+    ]
+  ]
+
   @schema_def ftyp: [
                 fields: [
                   major_brand: :str32,
@@ -197,6 +211,7 @@ defmodule Membrane.MP4.Container.Schema.Default do
                           avc3: @avc_schema,
                           hvc1: @hevc_schema,
                           hev1: @hevc_schema,
+                          av01: @av1_schema,
                           mp4a: [
                             fields: [
                               reserved: <<0::6*8>>,

lib/membrane_mp4/demuxer/isom.ex

@@ -50,7 +50,8 @@ defmodule Membrane.MP4.Demuxer.ISOM do
           stream_structure: {_hevc, _dcr},
           alignment: :au
         },
-        %Membrane.Opus{self_delimiting?: false}
+        %Membrane.Opus{self_delimiting?: false},
+        %Membrane.AV1{alignment: :tu}
       ),
     availability: :on_request,
     options: [
@@ -517,6 +518,7 @@ defmodule Membrane.MP4.Demuxer.ISOM do
 
   defp sample_description_to_kind(%Membrane.H264{}), do: :video
   defp sample_description_to_kind(%Membrane.H265{}), do: :video
+  defp sample_description_to_kind(%Membrane.AV1{}), do: :video
   defp sample_description_to_kind(%Membrane.AAC{}), do: :audio
   defp sample_description_to_kind(%Membrane.Opus{}), do: :audio
 

lib/membrane_mp4/demuxer/isom/samples_info.ex

@@ -5,11 +5,13 @@ defmodule Membrane.MP4.Demuxer.ISOM.SamplesInfo do
   # generating output buffers using that structure.
   # The samples' description is generated from the the `moov` box, which describes how the data is stored inside the `mdat` box.
 
-  alias Membrane.{Buffer, Time}
+  alias Membrane.{AV1, Buffer, Time}
   alias Membrane.MP4.Container
   alias Membrane.MP4.MovieBox.SampleTableBox
   alias Membrane.MP4.Track.SampleTable
 
+  @av1_temporal_delimiter_obu <<0x12, 0x00>>
+
   @enforce_keys [
     :samples,
     :tracks_number,
@@ -77,7 +79,7 @@ defmodule Membrane.MP4.Demuxer.ISOM.SamplesInfo do
 
         buffer =
           {%Buffer{
-             payload: payload,
+             payload: maybe_prepend_temporal_delimiter(samples_info, track_id, payload),
              dts: dts,
              pts: pts
            }, track_id}
@@ -95,6 +97,16 @@ defmodule Membrane.MP4.Demuxer.ISOM.SamplesInfo do
     end
   end
 
+  # The AV1 ISOBMFF binding requires samples to be stored without temporal delimiter OBUs,
+  # while `Membrane.AV1` with `alignment: :tu` requires each buffer to start with one,
+  # so it is restored here.
+  defp maybe_prepend_temporal_delimiter(samples_info, track_id, payload) do
+    case samples_info.sample_tables[track_id].sample_description do
+      %AV1{} -> @av1_temporal_delimiter_obu <> payload
+      _other -> payload
+    end
+  end
+
   defp get_dts_and_pts(samples_info, %{
          sample_delta: delta,
          track_id: track_id,

lib/membrane_mp4/helper.ex

@@ -17,5 +17,6 @@ defmodule Membrane.MP4.Helper do
   @spec key_frame?(Membrane.Buffer.metadata()) :: boolean()
   def key_frame?(%{h264: %{key_frame?: false}}), do: false
   def key_frame?(%{h265: %{key_frame?: false}}), do: false
+  def key_frame?(%{av1: %{key_frame?: false}}), do: false
   def key_frame?(_metadata), do: true
 end

lib/membrane_mp4/movie_box/sample_table_box.ex

@@ -3,7 +3,7 @@ defmodule Membrane.MP4.MovieBox.SampleTableBox do
 
   require Membrane.H264
   require Membrane.Logger
-  alias Membrane.{AAC, H264, H265, Opus}
+  alias Membrane.{AAC, AV1, H264, H265, Opus}
   alias Membrane.MP4.{Container, Helper, Track.SampleTable}
 
   @spec assemble(SampleTable.t()) :: Container.t()
@@ -184,6 +184,79 @@ defmodule Membrane.MP4.MovieBox.SampleTableBox do
     ]
   end
 
+  defp assemble_sample_description(%AV1{width: width, height: height} = format) do
+    av1c_content = build_av1c(format)
+
+    [
+      {:av01,
+       %{
+         children: [
+           av1C: %{
+             content: av1c_content
+           },
+           pasp: %{
+             children: [],
+             fields: %{h_spacing: 1, v_spacing: 1}
+           }
+         ],
+         fields: %{
+           compressor_name: <<0::size(32)-unit(8)>>,
+           depth: 24,
+           flags: 0,
+           frame_count: 1,
+           height: height || 0,
+           horizresolution: {0, 0},
+           num_of_entries: 1,
+           version: 0,
+           vertresolution: {0, 0},
+           width: width || 0
+         }
+       }}
+    ]
+  end
+
+  # Build av1C (AV1CodecConfigurationRecord) per ISO/IEC 14496-15
+  defp build_av1c(%AV1{} = format) do
+    config = av1c_config(format)
+
+    <<
+      # Byte 0: marker (1) | version (7)
+      1::1,
+      1::7,
+      # Byte 1: seq_profile (3) | seq_level_idx_0 (5)
+      config.profile::3,
+      config.level_idx::5,
+      # Byte 2: seq_tier_0 (1) | high_bitdepth (1) | twelve_bit (1) | monochrome (1) |
+      #         chroma_subsampling_x (1) | chroma_subsampling_y (1) | chroma_sample_position (2)
+      config.tier::1,
+      config.high_bitdepth::1,
+      config.twelve_bit::1,
+      config.monochrome::1,
+      config.chroma_subsampling_x::1,
+      config.chroma_subsampling_y::1,
+      config.chroma_sample_position::2,
+      # Byte 3: reserved (3) | initial_presentation_delay_present (1) | reserved (4)
+      0::3,
+      0::1,
+      0::4
+    >>
+  end
+
+  defp av1c_config(%AV1{} = format) do
+    %{
+      profile: AV1.profile_to_seq_profile(format.profile || :main),
+      level_idx: AV1.level_to_seq_level_idx(format.level || :"4.0"),
+      tier: AV1.tier_to_seq_tier(format.tier || :main),
+      # Defaults for 8-bit, 4:2:0
+      high_bitdepth: 0,
+      twelve_bit: 0,
+      monochrome: 0,
+      chroma_subsampling_x: 1,
+      chroma_subsampling_y: 1,
+      chroma_sample_position: 0
+    }
+  end
+
   defp assemble_sample_deltas(%{timescale: timescale, decoding_deltas: decoding_deltas}),
     do:
       Enum.map(decoding_deltas, fn %{sample_count: count, sample_delta: delta} ->
@@ -343,6 +416,20 @@ defmodule Membrane.MP4.MovieBox.SampleTableBox do
     %Opus{channels: boxes[:dOps].fields.output_channel_count, self_delimiting?: false}
   end
 
+  defp unpack_sample_description(%{children: [{:av01, %{children: boxes, fields: fields}}]}) do
+    <<_marker::1, _version::7, seq_profile::3, seq_level_idx::5, seq_tier::1, _rest::bitstring>> =
+      boxes[:av1C].content
+
+    %AV1{
+      alignment: :tu,
+      width: fields.width,
+      height: fields.height,
+      profile: AV1.seq_profile_to_profile(seq_profile),
+      level: AV1.seq_level_idx_to_level(seq_level_idx),
+      tier: AV1.seq_tier_to_tier(seq_tier)
+    }
+  end
+
   defp unpack_sample_description(%{children: [{sample_type, _sample_metadata}]}) do
     Membrane.Logger.warning("Unknown sample type: #{inspect(sample_type)}")
     nil

lib/membrane_mp4/movie_box/track_box.ex

@@ -13,7 +13,8 @@ defmodule Membrane.MP4.MovieBox.TrackBox do
 
   defguardp is_audio(track)
             when not is_struct(track.stream_format, Membrane.H264) and
-                   not is_struct(track.stream_format, Membrane.H265)
+                   not is_struct(track.stream_format, Membrane.H265) and
+                   not is_struct(track.stream_format, Membrane.AV1)
 
   @spec assemble(Track.t()) :: Container.t()
   def assemble(track) do

lib/membrane_mp4/muxer/isom.ex

@@ -24,7 +24,8 @@ defmodule Membrane.MP4.Muxer.ISOM do
           stream_structure: {_hevc, _dcr},
           alignment: :au
         },
-        %Membrane.Opus{self_delimiting?: false}
+        %Membrane.Opus{self_delimiting?: false},
+        %Membrane.AV1{alignment: :tu}
       ),
     availability: :on_request
 
@@ -182,11 +183,23 @@ defmodule Membrane.MP4.Muxer.ISOM do
     {[event: {:output, event}], state}
   end
 
-  defp handle_queue_item({Pad.ref(:input, pad_ref), {:buffer, buffer}}, state) do
+  defp handle_queue_item({Pad.ref(:input, pad_ref), {:buffer, %Buffer{} = buffer}}, state) do
     # In case DTS is not set, use PTS. This is the case for audio tracks or H264 originated
     # from an RTP stream. ISO base media file format specification uses DTS for calculating
     # decoding deltas, and so is the implementation of sample table in this plugin.
-    buffer = %{buffer | dts: Buffer.get_dts_or_pts(buffer)}
+    buffer = %Buffer{buffer | dts: Buffer.get_dts_or_pts(buffer)}
+
+    # For AV1 tracks, strip temporal delimiter OBUs from the payload
+    track = get_in(state, [:pad_to_track, pad_ref])
+
+    buffer =
+      case track.stream_format do
+        %Membrane.AV1{} ->
+          %Buffer{buffer | payload: strip_av1_temporal_delimiter_obu(buffer.payload)}
+
+        _other ->
+          buffer
+      end
 
     state
     |> update_in([:pad_to_track, pad_ref], &Track.store_sample(&1, buffer))
@@ -286,4 +299,11 @@ defmodule Membrane.MP4.Muxer.ISOM do
   defp shift_left([]), do: []
 
   defp shift_left([first | rest]), do: rest ++ [first]
+
+  # Strip temporal delimiter OBUs from AV1 payload
+  # Temporal delimiter OBU: header=0x12 (type=2, has_size=1), size=0x00
+  defp strip_av1_temporal_delimiter_obu(<<0x12, 0x00, rest::binary>>),
+    do: strip_av1_temporal_delimiter_obu(rest)
+
+  defp strip_av1_temporal_delimiter_obu(payload), do: payload
 end

lib/membrane_mp4/track.ex

@@ -8,7 +8,7 @@ defmodule Membrane.MP4.Track do
   """
   require Membrane.{H264, H265}
   alias __MODULE__.SampleTable
-  alias Membrane.{AAC, H264, H265}
+  alias Membrane.{AAC, AV1, H264, H265}
   alias Membrane.MP4.Helper
 
   @type t :: %__MODULE__{
@@ -90,6 +90,8 @@ defmodule Membrane.MP4.Track do
           | {:avc1, %{profile: binary(), compatibility: binary(), level: binary()}}
           | {:hvc1,
              %{profile: non_neg_integer(), tier: non_neg_integer(), level: non_neg_integer()}}
+          | {:av01,
+             %{profile: non_neg_integer(), level: non_neg_integer(), tier: non_neg_integer()}}
           | nil
 
   def get_encoding_info(%__MODULE__{
@@ -139,6 +141,16 @@ defmodule Membrane.MP4.Track do
     {hevc, map}
   end
 
+  def get_encoding_info(%__MODULE__{stream_format: %AV1{} = format}) do
+    map = %{
+      profile: AV1.profile_to_seq_profile(format.profile || :main),
+      level: AV1.level_to_seq_level_idx(format.level || :"4.0"),
+      tier: AV1.tier_to_seq_tier(format.tier || :main)
+    }
+
+    {:av01, map}
+  end
+
   def get_encoding_info(_unknown), do: nil
 
   defp get_timescale(stream_format) do
@@ -149,13 +161,13 @@ defmodule Membrane.MP4.Track do
       %Membrane.AAC{sample_rate: sample_rate} ->
         sample_rate
 
-      %module{framerate: nil} when module in [H264, H265] ->
+      %module{framerate: nil} when module in [AV1, H264, H265] ->
         30 * 1024
 
-      %module{framerate: {0, _denominator}} when module in [H264, H265] ->
+      %module{framerate: {0, _denominator}} when module in [AV1, H264, H265] ->
         30 * 1024
 
-      %module{framerate: {nominator, _denominator}} when module in [H264, H265] ->
+      %module{framerate: {nominator, _denominator}} when module in [AV1, H264, H265] ->
         nominator * 1024
     end
   end

lib/membrane_mp4/track/sample_table.ex

@@ -177,6 +177,11 @@ defmodule Membrane.MP4.Track.SampleTable do
     Map.update!(sample_table, :sync_samples, &[sample_table.sample_count | &1])
   end
 
+  defp maybe_store_sync_sample(sample_table, %Buffer{metadata: %{av1: av1}})
+       when av1.key_frame? do
+    Map.update!(sample_table, :sync_samples, &[sample_table.sample_count | &1])
+  end
+
   defp maybe_store_sync_sample(sample_table, _buffer), do: sample_table
 
   defp store_last_dts(sample_table, %Buffer{dts: dts}),

mix.exs

@@ -45,6 +45,7 @@ defmodule Membrane.MP4.Plugin.MixProject do
       {:membrane_h264_format, "~> 0.6.1"},
       {:membrane_h265_format, "~> 0.2.0"},
       {:membrane_opus_format, "~> 0.3.0"},
+      {:membrane_av1_format, "~> 0.3.1"},
       {:membrane_file_plugin, "~> 0.17.0"},
       {:membrane_timestamp_queue, "~> 0.2.1"},
       {:bunch, "~> 1.5"},
@@ -53,6 +54,7 @@ defmodule Membrane.MP4.Plugin.MixProject do
       {:membrane_opus_plugin, "~> 0.19.0", only: :test},
       {:membrane_stream_plugin, "~> 0.4.0", only: :test},
       {:membrane_fake_plugin, "~> 0.11.0", only: :test},
+      {:membrane_ivf_plugin, "~> 0.9.2", only: :test},
       {:ex_doc, ">= 0.40.0", only: :dev, runtime: false},
       {:dialyxir, ">= 0.0.0", only: :dev, runtime: false},
       {:credo, ">= 0.0.0", only: :dev, runtime: false}