Generate best effort timestamps (#7)

lib/membrane_h265_plugin/parser.ex

@@ -36,6 +36,7 @@ defmodule Membrane.H265.Parser do
 
   alias __MODULE__.{
     AUSplitter,
+    AUTimestampGenerator,
     DecoderConfigurationRecord,
     Format,
     NALuParser,
@@ -95,15 +96,6 @@ defmodule Membrane.H265.Parser do
                 be provided via this option (only available for `:annexb` output stream format).
                 """
               ],
-              framerate: [
-                spec: {pos_integer(), pos_integer()} | nil,
-                default: nil,
-                description: """
-                Framerate of the video, represented as a tuple consisting of a numerator and the
-                denominator.
-                Its value will be sent inside the output `Membrane.H265` stream format.
-                """
-              ],
               skip_until_keyframe: [
                 spec: boolean(),
                 default: true,
@@ -151,6 +143,35 @@ defmodule Membrane.H265.Parser do
                 transported in the DCR, when in hev1 they will be present only in the stream (in-band).
                 If not provided or set to nil the stream's structure will remain unchanged.
                 """
+              ],
+              generate_best_effort_timestamps: [
+                spec:
+                  false
+                  | %{
+                      :framerate => {pos_integer(), pos_integer()},
+                      optional(:add_dts_offset) => boolean()
+                    },
+                default: false,
+                description: """
+                Generates timestamps based on given `framerate`.
+
+                This option works only when `Membrane.RemoteStream` format arrives.
+
+                Keep in mind that the generated timestamps may be inaccurate and lead
+                to video getting out of sync with other media, therefore h265 stream
+                should be kept in a container that stores the timestamps alongside.
+
+                By default, the parser adds negative DTS offset to the timestamps,
+                so that in case of frame reorder (which always happens when B frames
+                are present) the DTS was always bigger than PTS. If that is not desired,
+                you can set `add_dts_offset: false`.
+
+                The calculated DTS/PTS may be wrong since we base it on access units' POC (Picture Order Count).
+                We assume here that the POC is continuous on a CVS (Coded Video Sequence) which is
+                not guaranteed by the H265 specification. An example where POC values may not be continuous is
+                when generating sub-bitstream from the main stream by deleting access units belonging to a
+                higher temporal sub-layer.
+                """
               ]
 
   @impl true
@@ -162,6 +183,14 @@ defmodule Membrane.H265.Parser do
         stream_structure -> stream_structure
       end
 
+    {au_timestamp_generator, framerate} =
+      if opts.generate_best_effort_timestamps do
+        {AUTimestampGenerator.new(opts.generate_best_effort_timestamps),
+         opts.generate_best_effort_timestamps.framerate}
+      else
+        {nil, nil}
+      end
+
     state = %{
       nalu_splitter: nil,
       nalu_parser: nil,
@@ -170,7 +199,7 @@ defmodule Membrane.H265.Parser do
       profile: nil,
       previous_buffer_timestamps: {nil, nil},
       output_alignment: opts.output_alignment,
-      framerate: opts.framerate,
+      framerate: framerate,
       skip_until_keyframe: opts.skip_until_keyframe,
       repeat_parameter_sets: opts.repeat_parameter_sets,
       frame_prefix: <<>>,
@@ -181,7 +210,8 @@ defmodule Membrane.H265.Parser do
       initial_spss: opts.spss,
       initial_ppss: opts.ppss,
       input_stream_structure: nil,
-      output_stream_structure: output_stream_structure
+      output_stream_structure: output_stream_structure,
+      au_timestamp_generator: au_timestamp_generator
     }
 
     {[], state}
@@ -509,7 +539,7 @@ defmodule Membrane.H265.Parser do
     Enum.flat_map_reduce(aus, state, fn au, state ->
       {au, stream_format_actions, state} = process_au_parameter_sets(au, ctx, state)
 
-      {pts, dts} = hd(au).timestamps
+      {{pts, dts}, state} = prepare_timestamps(au, state)
 
       buffers_actions = [
         buffer:
@@ -521,6 +551,22 @@ defmodule Membrane.H265.Parser do
     end)
   end
 
+  @spec prepare_timestamps(AUSplitter.access_unit(), state()) ::
+          {{Membrane.Time.t(), Membrane.Time.t()}, state()}
+  defp prepare_timestamps(au, state) do
+    if state.mode == :bytestream and state.au_timestamp_generator do
+      {timestamps, timestamp_generator} =
+        AUTimestampGenerator.generate_ts_with_constant_framerate(
+          au,
+          state.au_timestamp_generator
+        )
+
+      {timestamps, %{state | au_timestamp_generator: timestamp_generator}}
+    else
+      {Enum.find(au, &(&1.type in NALuTypes.vcl_nalu_types())).timestamps, state}
+    end
+  end
+
   @spec maybe_add_parameter_sets(AUSplitter.access_unit(), state()) :: AUSplitter.access_unit()
   defp maybe_add_parameter_sets(au, %{repeat_parameter_sets: false}), do: au
 

lib/membrane_h265_plugin/parser/au_timestamp_generator.ex

@@ -0,0 +1,111 @@
+defmodule Membrane.H265.Parser.AUTimestampGenerator do
+  @moduledoc false
+
+  require Membrane.H265.Parser.NALuTypes, as: NALuTypes
+
+  alias Membrane.H265.Parser.NALu
+
+  @type framerate :: {frames :: pos_integer(), seconds :: pos_integer()}
+
+  @type t :: %{
+          framerate: framerate,
+          max_frame_reorder: 0..15,
+          au_counter: non_neg_integer(),
+          key_frame_au_idx: non_neg_integer(),
+          prev_pic_order_cnt_lsb: integer(),
+          prev_pic_order_cnt_msb: integer()
+        }
+
+  @spec new(config :: %{:framerate => framerate, optional(:add_dts_offset) => boolean()}) :: t
+  def new(config) do
+    # To make sure that PTS >= DTS at all times, we take maximal possible
+    # frame reorder (which is 15 according to the spec) and subtract
+    # `max_frame_reorder * frame_duration` from each frame's DTS.
+    # This behaviour can be disabled by setting `add_dts_offset: false`.
+    max_frame_reorder = if Map.get(config, :add_dts_offset, true), do: 15, else: 0
+
+    %{
+      framerate: config.framerate,
+      max_frame_reorder: max_frame_reorder,
+      au_counter: 0,
+      key_frame_au_idx: 0,
+      prev_pic_order_cnt_lsb: 0,
+      prev_pic_order_cnt_msb: 0
+    }
+  end
+
+  @spec generate_ts_with_constant_framerate([NALu.t()], t) ::
+          {{pts :: non_neg_integer(), dts :: non_neg_integer()}, t}
+  def generate_ts_with_constant_framerate(au, state) do
+    %{
+      au_counter: au_counter,
+      key_frame_au_idx: key_frame_au_idx,
+      max_frame_reorder: max_frame_reorder,
+      framerate: {frames, seconds}
+    } = state
+
+    first_vcl_nalu = Enum.find(au, &(&1.type in NALuTypes.vcl_nalu_types()))
+    {poc, state} = calculate_poc(first_vcl_nalu, state)
+
+    key_frame_au_idx = if poc == 0, do: au_counter, else: key_frame_au_idx
+    pts = div((key_frame_au_idx + poc) * seconds * Membrane.Time.second(), frames)
+    dts = div((au_counter - max_frame_reorder) * seconds * Membrane.Time.second(), frames)
+
+    state = %{
+      state
+      | au_counter: au_counter + 1,
+        key_frame_au_idx: key_frame_au_idx
+    }
+
+    {{pts, dts}, state}
+  end
+
+  # Calculate picture order count according to section 8.3.1 of the ITU-T H265 specification
+  defp calculate_poc(vcl_nalu, state) do
+    max_pic_order_cnt_lsb = 2 ** (vcl_nalu.parsed_fields.log2_max_pic_order_cnt_lsb_minus4 + 4)
+
+    # We exclude CRA pictures from IRAP pictures since we have no way
+    # to assert the value of the flag NoRaslOutputFlag.
+    # If the CRA is the first access unit in the bytestream, the flag would be
+    # equal to 1 which reset the POC counter, and that condition is
+    # satisfied here since the initial value for prev_pic_order_cnt_msb and
+    # prev_pic_order_cnt_lsb are 0
+    {prev_pic_order_cnt_msb, prev_pic_order_cnt_lsb} =
+      if vcl_nalu.parsed_fields.nal_unit_type in 16..20 do
+        {0, 0}
+      else
+        {state.prev_pic_order_cnt_msb, state.prev_pic_order_cnt_lsb}
+      end
+
+    pic_order_cnt_lsb = vcl_nalu.parsed_fields.pic_order_cnt_lsb
+
+    pic_order_cnt_msb =
+      cond do
+        pic_order_cnt_lsb < prev_pic_order_cnt_lsb and
+            prev_pic_order_cnt_lsb - pic_order_cnt_lsb >= div(max_pic_order_cnt_lsb, 2) ->
+          prev_pic_order_cnt_msb + max_pic_order_cnt_lsb
+
+        pic_order_cnt_lsb > prev_pic_order_cnt_lsb and
+            pic_order_cnt_lsb - prev_pic_order_cnt_lsb > div(max_pic_order_cnt_lsb, 2) ->
+          prev_pic_order_cnt_msb - max_pic_order_cnt_lsb
+
+        true ->
+          prev_pic_order_cnt_msb
+      end
+
+    {prev_pic_order_cnt_lsb, prev_pic_order_cnt_msb} =
+      if vcl_nalu.type in [:radl_r, :radl_n, :rasl_r, :rasl_n] or
+           vcl_nalu.parsed_fields.nal_unit_type in 0..15//2 do
+        {prev_pic_order_cnt_lsb, prev_pic_order_cnt_msb}
+      else
+        {pic_order_cnt_lsb, pic_order_cnt_msb}
+      end
+
+    {pic_order_cnt_msb + pic_order_cnt_lsb,
+     %{
+       state
+       | prev_pic_order_cnt_lsb: prev_pic_order_cnt_lsb,
+         prev_pic_order_cnt_msb: prev_pic_order_cnt_msb
+     }}
+  end
+end

lib/membrane_h265_plugin/parser/nalu_parser.ex

@@ -71,25 +71,30 @@ defmodule Membrane.H265.Parser.NALuParser do
 
     type = NALuTypes.get_type(parsed_fields.nal_unit_type)
 
-    {parsed_fields, scheme_parser_state} =
-      parse_proper_nalu_type(nalu_body, scheme_parser_state, type)
-
-    # Mark nalu as invalid if there's no parameter sets
-    nalu_status =
-      if type in NALuTypes.vcl_nalu_types() and
-           not Map.has_key?(parsed_fields, :separate_colour_plane_flag),
-         do: :error,
-         else: :valid
-
-    nalu =
-      %NALu{
-        parsed_fields: parsed_fields,
-        type: type,
-        status: nalu_status,
-        stripped_prefix: prefix,
-        payload: unprefixed_nalu_payload,
-        timestamps: timestamps
-      }
+    {nalu, scheme_parser_state} =
+      try do
+        {parsed_fields, scheme_parser_state} =
+          parse_proper_nalu_type(nalu_body, scheme_parser_state, type)
+
+        {%NALu{
+           parsed_fields: parsed_fields,
+           type: type,
+           status: :valid,
+           stripped_prefix: prefix,
+           payload: unprefixed_nalu_payload,
+           timestamps: timestamps
+         }, scheme_parser_state}
+      catch
+        scheme_parser ->
+          {%NALu{
+             parsed_fields: SchemeParser.get_local_state(scheme_parser),
+             type: type,
+             status: :error,
+             stripped_prefix: prefix,
+             payload: unprefixed_nalu_payload,
+             timestamps: timestamps
+           }, scheme_parser_state}
+      end
 
     state = %{state | scheme_parser_state: scheme_parser_state}
 

lib/membrane_h265_plugin/parser/nalu_parser/schemes/pps.ex

@@ -10,6 +10,9 @@ defmodule Membrane.H265.Parser.NALuParser.Schemes.PPS do
     do: [
       field: {:pic_parameter_set_id, :ue},
       field: {:seq_parameter_set_id, :ue},
+      field: {:dependent_slice_segments_enabled_flag, :u1},
+      field: {:output_flag_present_flag, :u1},
+      field: {:num_extra_slice_header_bits, :u3},
       save_state_as_global_state: {&{:pps, &1}, [:pic_parameter_set_id]}
     ]
 end

lib/membrane_h265_plugin/parser/nalu_parser/schemes/slice.ex

@@ -3,7 +3,13 @@ defmodule Membrane.H265.Parser.NALuParser.Schemes.Slice do
   @behaviour Membrane.H265.Parser.NALuParser.Scheme
 
   @impl true
-  def defaults(), do: []
+  def defaults(),
+    do: [
+      dependent_slice_segment_flag: 0,
+      no_output_of_prior_pics_flag: 0,
+      pic_output_flag: 1,
+      pic_order_cnt_lsb: 0
+    ]
 
   @impl true
   def scheme(),
@@ -14,7 +20,35 @@ defmodule Membrane.H265.Parser.NALuParser.Schemes.Slice do
         field: {:no_output_of_prior_pics_flag, :u1}
       },
       field: {:pic_parameter_set_id, :ue},
-      execute: &load_data_from_sps(&1, &2, &3)
+      execute: &load_data_from_sps/3,
+      if: {
+        {&(&1 == 0), [:first_slice_segment_in_pic_flag]},
+        if: {
+          {&(&1 == 1), [:dependent_slice_segments_enabled_flag]},
+          field: {:dependent_slice_segment_flag, :u1}
+        },
+        field: {:slice_segment_address, {:uv, & &1, [:segment_addr_length]}}
+      },
+      if: {
+        {&(&1 == 0), [:dependent_slice_segment_flag]},
+        for: {
+          [iterator: :j, from: 0, to: {&(&1 - 1), [:num_extra_slice_header_bits]}],
+          field: {:slice_reserved_flag, :u1}
+        },
+        field: {:slice_type, :ue},
+        if: {
+          {&(&1 == 1), [:output_flag_present_flag]},
+          field: {:pic_output_flag, :u1}
+        },
+        if: {
+          {&(&1 == 1), [:separate_colour_plane_flag]},
+          field: {:colour_plane_id, :u2}
+        },
+        if: {
+          {&(&1 != 19 and &1 != 20), [:nal_unit_type]},
+          field: {:pic_order_cnt_lsb, {:uv, &(&1 + 4), [:log2_max_pic_order_cnt_lsb_minus4]}}
+        }
+      }
     ]
 
   defp load_data_from_sps(payload, state, _iterators) do
@@ -31,13 +65,22 @@ defmodule Membrane.H265.Parser.NALuParser.Schemes.Slice do
           Map.get(sps, :separate_colour_plane_flag, 0)
         )
 
-      sps_fields = Map.take(sps, [:log2_max_pic_order_cnt_lsb_minus4])
-
+      sps_fields = Map.take(sps, [:segment_addr_length, :log2_max_pic_order_cnt_lsb_minus4])
       state = Map.update(state, :__local__, %{}, &Map.merge(&1, sps_fields))
 
+      # PPS fields
+      pps_fields =
+        Map.take(pps, [
+          :dependent_slice_segments_enabled_flag,
+          :output_flag_present_flag,
+          :num_extra_slice_header_bits
+        ])
+
+      state = Map.update(state, :__local__, %{}, &Map.merge(&1, pps_fields))
+
       {payload, state}
     else
-      _error -> {payload, state}
+      _error -> throw(state)
     end
   end
 end