Merge pull request #2 from membraneframework/develop

Add frame size calculation

lib/membrane_caps_video_raw.ex

@@ -2,6 +2,7 @@ defmodule Membrane.Caps.Video.Raw do
   @moduledoc """
   This module provides caps struct for raw video frames.
   """
+  require Integer
 
   @typedoc """
   Width of single frame in pixels.
@@ -40,5 +41,51 @@ defmodule Membrane.Caps.Video.Raw do
         }
 
   @enforce_keys [:width, :height, :framerate, :format, :aligned]
-  defstruct [:width, :height, :framerate, :format, :aligned]
+  defstruct @enforce_keys
+
+  @doc """
+  Simple wrapper over `frame_size/3`. Returns the size of raw video frame
+  in bytes for the given caps.
+  """
+  @spec frame_size(t()) :: Bunch.Type.try_t(pos_integer)
+  def frame_size(%__MODULE__{format: format, width: width, height: height}) do
+    frame_size(format, width, height)
+  end
+
+  @doc """
+  Returns the size of raw video frame in bytes (without padding).
+
+  It may result in error when dimensions don't fulfill requirements for the given format
+  (e.g. I420 requires both dimensions to be divisible by 2).
+  """
+  @spec frame_size(Raw.format_t(), Raw.width_t(), Raw.height()) :: Bunch.Type.try_t(pos_integer)
+  def frame_size(format, width, height)
+      when format in [:I420, :YV12, :NV12, :NV21] and Integer.is_even(width) and
+             Integer.is_even(height) do
+    # Subsampling by 2 in both dimensions
+    # Y = width * height
+    # V = U = (width / 2) * (height / 2)
+    {:ok, div(width * height * 3, 2)}
+  end
+
+  def frame_size(:I422, width, height) when Integer.is_even(width) do
+    # Subsampling by 2 in horizontal dimension
+    # Y = width * height
+    # V = U = (width / 2) * height
+    {:ok, width * height * 2}
+  end
+
+  def frame_size(format, width, height) when format in [:I444, :RGB] do
+    # No subsampling
+    {:ok, width * height * 3}
+  end
+
+  def frame_size(format, width, height) when format in [:AYUV, :RGBA, :BGRA] do
+    # No subsampling and added alpha channel
+    {:ok, width * height * 4}
+  end
+
+  def frame_size(_, _, _) do
+    {:error, :invalid_dims}
+  end
 end

test/membrane_caps_video_raw_test.exs

@@ -0,0 +1,74 @@
+defmodule Membrane.Caps.Membrane.Caps.Video.Raw.Test do
+  use ExUnit.Case, async: true
+  @module Membrane.Caps.Video.Raw
+
+  test "frame_size for :I420 format" do
+    format = :I420
+    assert @module.frame_size(format, 10, 20) == {:ok, 300}
+    assert @module.frame_size(format, 9, 20) == {:error, :invalid_dims}
+    assert @module.frame_size(format, 10, 19) == {:error, :invalid_dims}
+  end
+
+  test "frame_size for :I422 format" do
+    format = :I422
+    assert @module.frame_size(format, 10, 20) == {:ok, 400}
+    assert @module.frame_size(format, 9, 20) == {:error, :invalid_dims}
+    assert @module.frame_size(format, 10, 19) == {:ok, 380}
+  end
+
+  test "frame_size for :I444 format" do
+    format = :I444
+    assert @module.frame_size(format, 10, 20) == {:ok, 600}
+    assert @module.frame_size(format, 9, 20) == {:ok, 540}
+    assert @module.frame_size(format, 10, 19) == {:ok, 570}
+  end
+
+  test "frame_size for :RGB format" do
+    format = :RGB
+    assert @module.frame_size(format, 10, 20) == {:ok, 600}
+    assert @module.frame_size(format, 9, 20) == {:ok, 540}
+    assert @module.frame_size(format, 10, 19) == {:ok, 570}
+  end
+
+  test "frame_size for :BGRA format" do
+    format = :BGRA
+    assert @module.frame_size(format, 10, 20) == {:ok, 800}
+    assert @module.frame_size(format, 9, 20) == {:ok, 720}
+    assert @module.frame_size(format, 10, 19) == {:ok, 760}
+  end
+
+  test "frame_size for :RGBA format" do
+    format = :RGBA
+    assert @module.frame_size(format, 10, 20) == {:ok, 800}
+    assert @module.frame_size(format, 9, 20) == {:ok, 720}
+    assert @module.frame_size(format, 10, 19) == {:ok, 760}
+  end
+
+  test "frame_size for :NV12 format" do
+    format = :NV12
+    assert @module.frame_size(format, 10, 20) == {:ok, 300}
+    assert @module.frame_size(format, 9, 20) == {:error, :invalid_dims}
+    assert @module.frame_size(format, 10, 19) == {:error, :invalid_dims}
+  end
+
+  test "frame_size for :NV21 format" do
+    format = :NV21
+    assert @module.frame_size(format, 10, 20) == {:ok, 300}
+    assert @module.frame_size(format, 9, 20) == {:error, :invalid_dims}
+    assert @module.frame_size(format, 10, 19) == {:error, :invalid_dims}
+  end
+
+  test "frame_size for :YV12 format" do
+    format = :YV12
+    assert @module.frame_size(format, 10, 20) == {:ok, 300}
+    assert @module.frame_size(format, 9, 20) == {:error, :invalid_dims}
+    assert @module.frame_size(format, 10, 19) == {:error, :invalid_dims}
+  end
+
+  test "frame_size for :AYUV format" do
+    format = :AYUV
+    assert @module.frame_size(format, 10, 20) == {:ok, 800}
+    assert @module.frame_size(format, 9, 20) == {:ok, 720}
+    assert @module.frame_size(format, 10, 19) == {:ok, 760}
+  end
+end