|
| 1 | +defmodule Membrane.VKVideo.Encoder do |
| 2 | + @moduledoc """ |
| 3 | + H.264 encoder taking advantage of hardware acceleration provided |
| 4 | + by Vulkan video extensions. |
| 5 | + """ |
| 6 | + use Membrane.Filter |
| 7 | + |
| 8 | + require Membrane.Logger |
| 9 | + |
| 10 | + alias Membrane.VKVideo.{DeviceServer, Native} |
| 11 | + |
| 12 | + def_input_pad :input, accepted_format: %Membrane.RawVideo{pixel_format: :NV12} |
| 13 | + |
| 14 | + def_output_pad :output, |
| 15 | + accepted_format: %Membrane.H264{stream_structure: :annexb, alignment: :au} |
| 16 | + |
| 17 | + def_options tune: [ |
| 18 | + spec: :low_latency | :high_quality, |
| 19 | + default: :low_latency, |
| 20 | + description: """ |
| 21 | + Specifies whether the encoder should be optimized for minimal latency (which is |
| 22 | + important in case of livestreams) or for higher quality (applicable to offline encoding). |
| 23 | + """ |
| 24 | + ], |
| 25 | + approx_framerate: [ |
| 26 | + spec: {non_neg_integer(), pos_integer()} | nil, |
| 27 | + default: nil, |
| 28 | + description: """ |
| 29 | + Framerate of the stream expressed in number of frames per second. |
| 30 | + It's only used by the rate control mechanism and therefore it does not need to be an exact |
| 31 | + value. If nil, the framerate will be read from the stream format's structure or set |
| 32 | + to fixed value of 30 frames per second if framerate is not provided by the stream format. |
| 33 | + """ |
| 34 | + ], |
| 35 | + rate_control: [ |
| 36 | + spec: |
| 37 | + :encoder_default |
| 38 | + | :disabled |
| 39 | + | {:variable_bitrate, __MODULE__.VariableBitrate.t()} |
| 40 | + | {:constant_bitrate, __MODULE__.ConstantBitrate.t()}, |
| 41 | + default: :encoder_default, |
| 42 | + description: """ |
| 43 | + Specifies which rate control mechanism should by used by the encoder. |
| 44 | + """ |
| 45 | + ] |
| 46 | + |
| 47 | + @impl true |
| 48 | + def handle_init(_ctx, opts) do |
| 49 | + state = %{ |
| 50 | + encoder: nil, |
| 51 | + width: nil, |
| 52 | + height: nil, |
| 53 | + override_framerate?: opts.approx_framerate != nil, |
| 54 | + framerate: opts.approx_framerate, |
| 55 | + rate_control: opts.rate_control, |
| 56 | + tune: opts.tune |
| 57 | + } |
| 58 | + |
| 59 | + {[], state} |
| 60 | + end |
| 61 | + |
| 62 | + @impl true |
| 63 | + def handle_stream_format(:input, stream_format, _ctx, state) do |
| 64 | + cond do |
| 65 | + state.override_framerate? and |
| 66 | + (stream_format.width != state.width or |
| 67 | + stream_format.height != state.height) -> |
| 68 | + %{ |
| 69 | + state |
| 70 | + | width: stream_format.width, |
| 71 | + height: stream_format.height |
| 72 | + } |
| 73 | + |> spawn_encoder() |
| 74 | + |
| 75 | + not state.override_framerate? and |
| 76 | + (stream_format.width != state.width or stream_format.height != state.height or |
| 77 | + stream_format.framerate != state.framerate) -> |
| 78 | + %{ |
| 79 | + state |
| 80 | + | width: stream_format.width, |
| 81 | + height: stream_format.height, |
| 82 | + framerate: resolve_framerate(stream_format, state) |
| 83 | + } |
| 84 | + |> spawn_encoder() |
| 85 | + |
| 86 | + true -> |
| 87 | + {[], state} |
| 88 | + end |
| 89 | + end |
| 90 | + |
| 91 | + defp spawn_encoder(state) do |
| 92 | + {:ok, device} = DeviceServer.get_device() |
| 93 | + |
| 94 | + {:ok, encoder} = |
| 95 | + Native.new_encoder( |
| 96 | + device, |
| 97 | + state.width, |
| 98 | + state.height, |
| 99 | + state.framerate, |
| 100 | + state.tune, |
| 101 | + state.rate_control |
| 102 | + ) |
| 103 | + |
| 104 | + state = put_in(state, [:encoder], encoder) |
| 105 | + |
| 106 | + stream_format = |
| 107 | + %Membrane.H264{ |
| 108 | + stream_structure: :annexb, |
| 109 | + alignment: :au, |
| 110 | + width: state.width, |
| 111 | + height: state.height |
| 112 | + } |
| 113 | + |
| 114 | + stream_format = |
| 115 | + if state.override_framerate? do |
| 116 | + stream_format |
| 117 | + else |
| 118 | + %{stream_format | framerate: state.framerate} |
| 119 | + end |
| 120 | + |
| 121 | + {[stream_format: {:output, stream_format}], state} |
| 122 | + end |
| 123 | + |
| 124 | + defp resolve_framerate(stream_format, state) do |
| 125 | + if is_nil(stream_format.framerate) and requires_framerate?(state.rate_control) do |
| 126 | + Membrane.Logger.warning(""" |
| 127 | + Framerate is required when using #{inspect(elem(state.rate_control, 0))} rate control but it |
| 128 | + wasn't provided in the stream format nor via options. Please provide approximate framerate |
| 129 | + using `approx_framerate` option of the element. |
| 130 | + """) |
| 131 | + end |
| 132 | + |
| 133 | + stream_format.framerate || {30, 1} |
| 134 | + end |
| 135 | + |
| 136 | + defp requires_framerate?({:constant_bitrate, _constant_bitrate}), do: true |
| 137 | + defp requires_framerate?({:variable_bitrate, _variable_bitrate}), do: true |
| 138 | + defp requires_framerate?(_other_rate_control), do: false |
| 139 | + |
| 140 | + @impl true |
| 141 | + def handle_buffer(:input, buffer, _ctx, state) do |
| 142 | + {:ok, encoded_frame} = Native.encode(state.encoder, buffer.payload, buffer.pts) |
| 143 | + |
| 144 | + pts = |
| 145 | + if encoded_frame.pts_ns != nil, |
| 146 | + do: Membrane.Time.nanoseconds(encoded_frame.pts_ns), |
| 147 | + else: nil |
| 148 | + |
| 149 | + {[ |
| 150 | + buffer: |
| 151 | + {:output, |
| 152 | + %Membrane.Buffer{ |
| 153 | + payload: encoded_frame.payload, |
| 154 | + pts: pts, |
| 155 | + dts: buffer.dts |
| 156 | + }} |
| 157 | + ], state} |
| 158 | + end |
| 159 | + |
| 160 | + @impl true |
| 161 | + def handle_end_of_stream(:input, _ctx, state) do |
| 162 | + :ok = Native.destroy(state.encoder) |
| 163 | + state = %{state | encoder: nil} |
| 164 | + {[end_of_stream: :output], state} |
| 165 | + end |
| 166 | +end |
0 commit comments