Forward timestamps

lib/membrane_ffmpeg_swresample_plugin/converter.ex

@@ -64,7 +64,8 @@ defmodule Membrane.FFmpeg.SWResample.Converter do
       |> Map.merge(%{
         native: nil,
         queue: <<>>,
-        input_stream_format_provided?: options.input_stream_format != nil
+        input_stream_format_provided?: options.input_stream_format != nil,
+        pts_queue: []
       })
 
     {[], state}
@@ -122,16 +123,29 @@ defmodule Membrane.FFmpeg.SWResample.Converter do
   end
 
   @impl true
-  def handle_buffer(:input, %Buffer{payload: payload}, _ctx, state) do
+  def handle_buffer(:input, %Buffer{payload: payload, pts: input_pts}, _ctx, state) do
+    input_frame_size = RawAudio.frame_size(state.input_stream_format)
+    output_frame_size = RawAudio.frame_size(state.output_stream_format)
+    ratio = output_frame_size / input_frame_size
+
+    expected_output_frames_count = (byte_size(payload) / input_frame_size * ratio) |> trunc()
+
+    state =
+      Map.update!(state, :pts_queue, fn pts_queue ->
+        pts_queue ++ [{input_pts, expected_output_frames_count}]
+      end)
+
     conversion_result =
-      convert!(state.native, RawAudio.frame_size(state.input_stream_format), payload, state.queue)
+      convert!(state.native, input_frame_size, payload, state.queue)
 
     case conversion_result do
       {<<>>, queue} ->
         {[], %{state | queue: queue}}
 
       {converted, queue} ->
-        {[buffer: {:output, %Buffer{payload: converted}}], %{state | queue: queue}}
+        {state, out_pts} = update_pts_queue(state, byte_size(converted) / output_frame_size)
+
+        {[buffer: {:output, %Buffer{payload: converted, pts: out_pts}}], %{state | queue: queue}}
     end
   end
 
@@ -154,8 +168,15 @@ defmodule Membrane.FFmpeg.SWResample.Converter do
         {[end_of_stream: :output], %{state | queue: <<>>}}
 
       converted ->
-        {[buffer: {:output, %Buffer{payload: converted}}, end_of_stream: :output],
-         %{state | queue: <<>>}}
+        converted_frames_count =
+          byte_size(converted) / RawAudio.frame_size(state.output_stream_format)
+
+        {state, out_pts} = update_pts_queue(state, converted_frames_count)
+
+        {[
+           buffer: {:output, %Buffer{payload: converted, pts: out_pts}},
+           end_of_stream: :output
+         ], %{state | queue: <<>>}}
     end
   end
 
@@ -212,4 +233,15 @@ defmodule Membrane.FFmpeg.SWResample.Converter do
       {:error, reason} -> raise "Error while flushing converter: #{inspect(reason)}"
     end
   end
+
+  defp update_pts_queue(state, converted_frames_count) do
+    {[{out_pts, expected_frames}], rest} = Enum.split(state.pts_queue, 1)
+
+    if converted_frames_count < expected_frames do
+      {%{state | pts_queue: [{out_pts, expected_frames - converted_frames_count}] ++ rest},
+       out_pts}
+    else
+      {%{state | pts_queue: rest}, out_pts}
+    end
+  end
 end

test/membrane_ffmpeg_swresample_plugin/converter_test.exs

@@ -27,7 +27,8 @@ defmodule Membrane.FFmpeg.SWResample.ConverterTest do
         output_stream_format: @u8_format,
         frames_per_buffer: 2048,
         native: nil,
-        queue: <<>>
+        queue: <<>>,
+        pts_queue: []
       }
     }
   end
@@ -156,7 +157,7 @@ defmodule Membrane.FFmpeg.SWResample.ConverterTest do
 
       assert {[], new_state} = @module.handle_buffer(:input, buffer, nil, state)
 
-      assert new_state == %{state | queue: payload}
+      assert %{new_state | pts_queue: nil} == %{state | queue: payload, pts_queue: nil}
       refute_called(@native, :convert)
     end
 

test/membrane_ffmpeg_swresample_plugin/pts_forward_test.exs

@@ -0,0 +1,64 @@
+defmodule Membrane.FFmpeg.SWResample.PtsForwardTest do
+  use ExUnit.Case
+
+  import Membrane.ChildrenSpec
+  import Membrane.Testing.Assertions
+
+  alias Membrane.FFmpeg.SWResample.Converter
+  alias Membrane.{RawAudio, Testing}
+
+  test "pts forward test" do
+    input_stream_format = %RawAudio{sample_format: :s16le, sample_rate: 16_000, channels: 2}
+    output_stream_format = %RawAudio{sample_format: :s32le, sample_rate: 32_000, channels: 2}
+
+    # 32 frames * 2048 bytes
+    fixture_path = Path.expand(Path.join(__DIR__, "/../fixtures/input_s16le_stereo_16khz.raw"))
+
+    spec = [
+      child(:source, %Membrane.Testing.Source{output: buffers_from_file(fixture_path)})
+      |> child(:resampler, %Converter{
+        input_stream_format: input_stream_format,
+        output_stream_format: output_stream_format
+      })
+      |> child(:sink, Membrane.Testing.Sink)
+    ]
+
+    pipeline = Testing.Pipeline.start_link_supervised!(spec: spec)
+    assert_start_of_stream(pipeline, :sink)
+    assert_sink_buffer(pipeline, :sink, _buffer)
+
+    Enum.each(0..30, fn index ->
+      assert_sink_buffer(pipeline, :sink, %Membrane.Buffer{pts: out_pts})
+      assert out_pts == index * 31_250_000
+    end)
+
+    assert_sink_buffer(pipeline, :sink, %Membrane.Buffer{pts: 937_500_000})
+
+    assert_end_of_stream(pipeline, :sink)
+    Testing.Pipeline.terminate(pipeline)
+  end
+
+  defp buffers_from_file(path) do
+    binary = File.read!(path)
+
+    split_binary(binary)
+    |> Enum.with_index()
+    |> Enum.map(fn {payload, index} ->
+      %Membrane.Buffer{
+        payload: payload,
+        pts: index * 31_250_000
+      }
+    end)
+  end
+
+  @spec split_binary(binary(), list(binary())) :: list(binary())
+  def split_binary(binary, acc \\ [])
+
+  def split_binary(<<binary::binary-size(2048), rest::binary>>, acc) do
+    split_binary(rest, [binary] ++ acc)
+  end
+
+  def split_binary(rest, acc) when byte_size(rest) <= 2048 do
+    Enum.reverse(acc) ++ [rest]
+  end
+end