defmodule Membrane.H265.AUSplitter do
@moduledoc """
Module providing functionalities to group H265 NAL units
into access units.
The access unit splitter's behaviour is based on section **7.4.2.4.4**
*"Order of NAL units and coded pictures and association to access units"*
of the *"ITU-T Rec. H.265 (08/2021)"* specification.
"""
@behaviour Membrane.H26x.AUSplitter
require Logger
require Membrane.H265.NALuTypes, as: NALuTypes
alias Membrane.H265.NALuTypes
alias Membrane.H26x.{AUSplitter, NALu}
@typedoc """
A structure holding a state of the access unit splitter.
"""
@opaque t :: %__MODULE__{
nalus_acc: [NALu.t()],
fsm_state: :first | :second,
previous_nalu: NALu.t() | nil,
access_units_to_output: [AUSplitter.access_unit()]
}
@enforce_keys [
:nalus_acc,
:fsm_state,
:previous_nalu,
:access_units_to_output
]
defstruct @enforce_keys
@doc """
Returns a structure holding a clear state of the
access unit splitter.
"""
@spec new() :: t()
def new() do
%__MODULE__{
nalus_acc: [],
fsm_state: :first,
previous_nalu: nil,
access_units_to_output: []
}
end
@non_vcl_nalus_at_au_beginning [:vps, :sps, :pps, :prefix_sei]
@non_vcl_nalus_at_au_end [:fd, :eos, :eob, :suffix_sei]
@doc """
Splits the given list of NAL units into the access units.
It can be used for a stream which is not completely available at the time of function invocation,
as the function updates the state of the access unit splitter - the function can
be invoked once more, with new NAL units and the updated state.
Under the hood, `split/2` defines a finite state machine
with two states: `:first` and `:second`. The state `:first` describes the state before
reaching the first segment of a coded picture NALu of a given access unit. The state `:second`
describes the state after processing the first segment of the coded picture of a given
access unit.
"""
@spec split([NALu.t()], boolean(), t()) :: {[AUSplitter.access_unit()], t()}
def split(nalus, assume_au_aligned \\ false, state) do
state = do_split(nalus, state)
{aus, state} =
if assume_au_aligned do
{state.access_units_to_output ++ [state.nalus_acc],
%__MODULE__{state | access_units_to_output: [], nalus_acc: []}}
else
{state.access_units_to_output, %__MODULE__{state | access_units_to_output: []}}
end
{Enum.reject(aus, &Enum.empty?/1), state}
end
defp do_split([first_nalu | rest_nalus], %{fsm_state: :first} = state) do
cond do
access_unit_first_slice_segment?(first_nalu) ->
do_split(
rest_nalus,
%__MODULE__{
state
| nalus_acc: state.nalus_acc ++ [first_nalu],
fsm_state: :second,
previous_nalu: first_nalu
}
)
(first_nalu.type == :aud and state.nalus_acc == []) or
first_nalu.type in @non_vcl_nalus_at_au_beginning or
NALu.int_type(first_nalu) in 41..44 or
NALu.int_type(first_nalu) in 48..55 ->
do_split(
rest_nalus,
%__MODULE__{state | nalus_acc: state.nalus_acc ++ [first_nalu]}
)
true ->
Logger.warning("AUSplitter: Improper transition")
do_split(rest_nalus, state)
end
end
defp do_split([first_nalu | rest_nalus], %{fsm_state: :second} = state) do
previous_nalu = state.previous_nalu
cond do
first_nalu.type == :aud or first_nalu.type in @non_vcl_nalus_at_au_beginning ->
do_split(
rest_nalus,
%__MODULE__{
state
| nalus_acc: [first_nalu],
fsm_state: :first,
access_units_to_output: state.access_units_to_output ++ [state.nalus_acc]
}
)
access_unit_first_slice_segment?(first_nalu) ->
do_split(
rest_nalus,
%__MODULE__{
state
| nalus_acc: [first_nalu],
previous_nalu: first_nalu,
access_units_to_output: state.access_units_to_output ++ [state.nalus_acc]
}
)
first_nalu.type == previous_nalu.type or
first_nalu.type in @non_vcl_nalus_at_au_end or
NALu.int_type(first_nalu) in 45..47 or
NALu.int_type(first_nalu) in 56..63 ->
do_split(
rest_nalus,
%__MODULE__{
state
| nalus_acc: state.nalus_acc ++ [first_nalu],
previous_nalu: first_nalu
}
)
true ->
Logger.warning("AUSplitter: Improper transition")
do_split(rest_nalus, state)
end
end
defp do_split([], state) do
state
end
defp access_unit_first_slice_segment?(nalu) do
NALuTypes.is_vcl_nalu_type(nalu.type) and
nalu.parsed_fields[:first_slice_segment_in_pic_flag] == 1
end
end