Add ability to specify output parameters

elixir-webrtc · Aug 7, 2024 · 23d3aaf · 23d3aaf
1 parent dad836a
commit 23d3aaf
Show file tree

Hide file tree

Showing 11 changed files with 229 additions and 36 deletions.
diff --git a/c_src/xav/xav_decoder.c b/c_src/xav/xav_decoder.c
@@ -7,7 +7,7 @@ ErlNifResourceType *xav_decoder_resource_type;
 static int init_audio_converter(struct XavDecoder *xav_decoder);
 
 ERL_NIF_TERM new(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
-  if (argc != 1) {
+  if (argc != 4) {
     return xav_nif_raise(env, "invalid_arg_count");
   }
 
@@ -22,10 +22,34 @@ ERL_NIF_TERM new(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
     return xav_nif_raise(env, "failed_to_get_atom");
   }
 
+  unsigned int out_format_len;
+  if (!enif_get_atom_length(env, argv[1], &out_format_len, ERL_NIF_LATIN1)) {
+    return xav_nif_raise(env, "failed_to_get_atom_length");
+  }
+
+  char *out_format = (char *)XAV_ALLOC((out_format_len + 1) * sizeof(char *));
+
+  if (enif_get_atom(env, argv[1], out_format, out_format_len + 1, ERL_NIF_LATIN1) == 0) {
+    return xav_nif_raise(env, "failed_to_get_atom");
+  }
+
+  int out_sample_rate;
+  if (!enif_get_int(env, argv[2], &out_sample_rate)) {
+    return xav_nif_raise(env, "invalid_out_sample_rate");
+  }
+
+  int out_channels;
+  if (!enif_get_int(env, argv[3], &out_channels)) {
+    return xav_nif_raise(env, "invalid_out_channels");
+  }
+
   struct XavDecoder *xav_decoder =
       enif_alloc_resource(xav_decoder_resource_type, sizeof(struct XavDecoder));
   xav_decoder->decoder = NULL;
   xav_decoder->ac = NULL;
+  xav_decoder->out_format = out_format;
+  xav_decoder->out_sample_rate = out_sample_rate;
+  xav_decoder->out_channels = out_channels;
 
   xav_decoder->decoder = decoder_alloc();
   if (xav_decoder->decoder == NULL) {
@@ -139,16 +163,47 @@ static int init_audio_converter(struct XavDecoder *xav_decoder) {
     return -1;
   }
 
-  int out_sample_rate = xav_decoder->decoder->c->sample_rate;
-  enum AVSampleFormat out_sample_fmt = AV_SAMPLE_FMT_FLT;
+  int out_sample_rate;
+  if (xav_decoder->out_sample_rate == 0) {
+    out_sample_rate = xav_decoder->decoder->c->sample_rate;
+  } else {
+    out_sample_rate = xav_decoder->out_sample_rate;
+  }
+
+  enum AVSampleFormat out_sample_fmt;
+  if (strcmp(xav_decoder->out_format, "u8") == 0) {
+    out_sample_fmt = AV_SAMPLE_FMT_U8;
+  } else if (strcmp(xav_decoder->out_format, "s16") == 0) {
+    out_sample_fmt = AV_SAMPLE_FMT_S16;
+  } else if (strcmp(xav_decoder->out_format, "s32") == 0) {
+    out_sample_fmt = AV_SAMPLE_FMT_S32;
+  } else if (strcmp(xav_decoder->out_format, "s64") == 0) {
+    out_sample_fmt = AV_SAMPLE_FMT_S64;
+  } else if (strcmp(xav_decoder->out_format, "f32") == 0) {
+    out_sample_fmt = AV_SAMPLE_FMT_FLT;
+  } else if (strcmp(xav_decoder->out_format, "f64") == 0) {
+    out_sample_fmt = AV_SAMPLE_FMT_DBL;
+  } else if (strcmp(xav_decoder->out_format, "nil") == 0) {
+    out_sample_fmt = av_get_alt_sample_fmt(xav_decoder->decoder->c->sample_fmt, 0);
+  } else {
+    return -1;
+  }
 
   struct ChannelLayout in_chlayout, out_chlayout;
 #if LIBAVUTIL_VERSION_MAJOR >= 58
   in_chlayout.layout = xav_decoder->decoder->c->ch_layout;
-  out_chlayout.layout = xav_decoder->decoder->c->ch_layout;
+  if (xav_decoder->out_channels == 0) {
+    out_chlayout.layout = in_chlayout.layout;
+  } else {
+    av_channel_layout_default(&out_chlayout.layout, xav_decoder->out_channels);
+  }
 #else
   in_chlayout.layout = xav_decoder->decoder->c->channel_layout;
-  out_chlayout.layout = xav_decoder->decoder->c->channel_layout;
+  if (xav_decoder->out_channels == 0) {
+    out_chlayout.layout = in_chlayout.layout;
+  } else {
+    out_chlayout.layout = av_get_default_channel_layout(xav_decoder->out_channels);
+  }
 #endif
 
   return audio_converter_init(xav_decoder->ac, in_chlayout, xav_decoder->decoder->c->sample_rate,
@@ -168,7 +223,7 @@ void free_xav_decoder(ErlNifEnv *env, void *obj) {
   }
 }
 
-static ErlNifFunc xav_funcs[] = {{"new", 1, new},
+static ErlNifFunc xav_funcs[] = {{"new", 4, new},
                                  {"decode", 4, decode, ERL_NIF_DIRTY_JOB_CPU_BOUND}};
 
 static int load(ErlNifEnv *env, void **priv, ERL_NIF_TERM load_info) {

diff --git a/c_src/xav/xav_decoder.h b/c_src/xav/xav_decoder.h
@@ -4,4 +4,7 @@
 struct XavDecoder {
   struct Decoder *decoder;
   struct AudioConverter *ac;
+  char *out_format;
+  int out_sample_rate;
+  int out_channels;
 };
diff --git a/c_src/xav/xav_reader.c b/c_src/xav/xav_reader.c
@@ -5,7 +5,7 @@ static int init_audio_converter(struct XavReader *xav_reader);
 ErlNifResourceType *xav_reader_resource_type;
 
 ERL_NIF_TERM new(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
-  if (argc != 3) {
+  if (argc != 6) {
     return xav_nif_raise(env, "invalid_arg_count");
   }
 
@@ -31,10 +31,34 @@ ERL_NIF_TERM new(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
     media_type = AVMEDIA_TYPE_AUDIO;
   }
 
+  unsigned int out_format_len;
+  if (!enif_get_atom_length(env, argv[3], &out_format_len, ERL_NIF_LATIN1)) {
+    return xav_nif_raise(env, "failed_to_get_atom_length");
+  }
+
+  char *out_format = (char *)XAV_ALLOC((out_format_len + 1) * sizeof(char *));
+
+  if (enif_get_atom(env, argv[3], out_format, out_format_len + 1, ERL_NIF_LATIN1) == 0) {
+    return xav_nif_raise(env, "failed_to_get_atom");
+  }
+
+  int out_sample_rate;
+  if (!enif_get_int(env, argv[4], &out_sample_rate)) {
+    return xav_nif_raise(env, "invalid_out_sample_rate");
+  }
+
+  int out_channels;
+  if (!enif_get_int(env, argv[5], &out_channels)) {
+    return xav_nif_raise(env, "invalid_out_channels");
+  }
+
   struct XavReader *xav_reader =
       enif_alloc_resource(xav_reader_resource_type, sizeof(struct XavReader));
   xav_reader->reader = NULL;
   xav_reader->ac = NULL;
+  xav_reader->out_format = out_format;
+  xav_reader->out_sample_rate = out_sample_rate;
+  xav_reader->out_channels = out_channels;
 
   xav_reader->reader = reader_alloc();
   if (xav_reader->reader == NULL) {
@@ -151,16 +175,42 @@ static int init_audio_converter(struct XavReader *xav_reader) {
     return -1;
   }
 
-  int out_sample_rate = 16000;
-  enum AVSampleFormat out_sample_fmt = AV_SAMPLE_FMT_FLT;
+  int out_sample_rate;
+  if (xav_reader->out_sample_rate == 0) {
+    out_sample_rate = xav_reader->reader->c->sample_rate;
+  } else {
+    out_sample_rate = xav_reader->out_sample_rate;
+  }
+
+  enum AVSampleFormat out_sample_fmt;
+  if (strcmp(xav_reader->out_format, "u8") == 0) {
+    out_sample_fmt = AV_SAMPLE_FMT_U8;
+  } else if (strcmp(xav_reader->out_format, "s16") == 0) {
+    out_sample_fmt = AV_SAMPLE_FMT_S16;
+  } else if (strcmp(xav_reader->out_format, "s32") == 0) {
+    out_sample_fmt = AV_SAMPLE_FMT_S32;
+  } else if (strcmp(xav_reader->out_format, "s64") == 0) {
+    out_sample_fmt = AV_SAMPLE_FMT_S64;
+  } else if (strcmp(xav_reader->out_format, "f32") == 0) {
+    out_sample_fmt = AV_SAMPLE_FMT_FLT;
+  } else if (strcmp(xav_reader->out_format, "f64") == 0) {
+    out_sample_fmt = AV_SAMPLE_FMT_DBL;
+  } else if (strcmp(xav_reader->out_format, "nil") == 0) {
+    out_sample_fmt = av_get_alt_sample_fmt(xav_reader->reader->c->sample_fmt, 0);
+  } else {
+    return -1;
+  }
 
   struct ChannelLayout in_chlayout, out_chlayout;
 #if LIBAVUTIL_VERSION_MAJOR >= 58
   in_chlayout.layout = xav_reader->reader->c->ch_layout;
-  av_channel_layout_from_mask(&out_chlayout.layout, AV_CH_LAYOUT_MONO);
+  if (xav_reader->out_channels == 0) {
+    out_chlayout.layout = in_chlayout.layout;
+  } else {
+    av_channel_layout_default(&out_chlayout.layout, xav_reader->out_channels);
+  }
 #else
   in_chlayout.layout = xav_reader->reader->c->channel_layout;
-  out_chlayout.layout = AV_CH_LAYOUT_MONO;
 
   if (xav_reader->reader->c->channel_layout == 0 && xav_reader->reader->c->channels > 0) {
     // In newer FFmpeg versions, 0 means that the order of channels is
@@ -176,6 +226,11 @@ static int init_audio_converter(struct XavReader *xav_reader) {
     return -1;
   }
 
+  if (xav_reader->out_channels == 0) {
+    out_chlayout.layout = in_chlayout.layout;
+  } else {
+    out_chlayout.layout = av_get_default_channel_layout(xav_reader->out_channels);
+  }
 #endif
 
   return audio_converter_init(xav_reader->ac, in_chlayout, xav_reader->reader->c->sample_rate,
@@ -195,7 +250,7 @@ void free_xav_reader(ErlNifEnv *env, void *obj) {
   }
 }
 
-static ErlNifFunc xav_funcs[] = {{"new", 3, new},
+static ErlNifFunc xav_funcs[] = {{"new", 6, new},
                                  {"next_frame", 1, next_frame, ERL_NIF_DIRTY_JOB_CPU_BOUND}};
 
 static int load(ErlNifEnv *env, void **priv, ERL_NIF_TERM load_info) {

diff --git a/c_src/xav/xav_reader.h b/c_src/xav/xav_reader.h
@@ -4,4 +4,7 @@
 struct XavReader {
   struct Reader *reader;
   struct AudioConverter *ac;
+  char *out_format;
+  int out_sample_rate;
+  int out_channels;
 };
diff --git a/lib/decoder.ex b/lib/decoder.ex
@@ -10,19 +10,52 @@ defmodule Xav.Decoder do
 
   @type t() :: reference()
 
+  @type opts :: [
+          out_format: Xav.Frame.format(),
+          out_sample_rate: integer(),
+          out_channels: integer()
+        ]
+
   @doc """
   Creates a new decoder.
+
+  `opts` can be used to specify desired output parameters.
+
+  E.g. if you want to change audio samples format just pass:
+
+  ```elixir
+    [out_format: :f32]
+  ```
+
+  Video frames are always returned in RGB format.
+  This setting cannot be changed.
+
+  Audio samples are always in the packed form -
+  samples from different channels are interleaved in the same, single binary:
+
+  <<c10, c20, c30, c11, c21, c31, c12, c22, c32>>
+
+  The way in which samples are interleaved is not specified.
+
+  An alternative would be to return a list of binaries, where
+  each binary represents different channel:
+
+  [
+    <<c10, c11, c12, c13, c14>>,
+    <<c20, c21, c22, c23, c24>>,
+    <<c30, c31, c32, c33, c34>>
+  ]
   """
-  @spec new(codec()) :: t()
-  def new(codec) do
-    Xav.Decoder.NIF.new(codec)
+  @spec new(codec(), opts()) :: t()
+  def new(codec, opts \\ []) do
+    out_format = opts[:out_format]
+    out_sample_rate = opts[:out_sample_rate] || 0
+    out_channels = opts[:out_channels] || 0
+    Xav.Decoder.NIF.new(codec, out_format, out_sample_rate, out_channels)
   end
 
   @doc """
-  Decodes an audio or video frame.
-
-  Video frames are always in the RGB format.
-  Audio samples are always interleaved.
+  Decodes an audio/video frame.
   """
   @spec decode(t(), binary(), pts: integer(), dts: integer()) ::
           {:ok, Xav.Frame.t()} | {:error, atom()}

diff --git a/lib/decoder_nif.ex b/lib/decoder_nif.ex
@@ -8,7 +8,7 @@ defmodule Xav.Decoder.NIF do
     :ok = :erlang.load_nif(path, 0)
   end
 
-  def new(_codec), do: :erlang.nif_error(:undef)
+  def new(_codec, _out_format, _out_sample_rate, _out_channels), do: :erlang.nif_error(:undef)
 
   def decode(_decoder, _data, _pts, _dts), do: :erlang.nif_error(:undef)
 end
diff --git a/lib/frame.ex b/lib/frame.ex
@@ -3,10 +3,14 @@ defmodule Xav.Frame do
   Audio/video frame.
   """
 
+  @type audio_format() :: :u8 | :s16 | :s32 | :s64 | :f32 | :f64
+  @type video_format() :: :rgb
+  @type format() :: audio_format() | video_format()
+
   @type t() :: %__MODULE__{
           type: :audio | :video,
           data: binary(),
-          format: atom(),
+          format: format(),
           width: non_neg_integer() | nil,
           height: non_neg_integer() | nil,
           samples: integer() | nil,
@@ -23,7 +27,10 @@ defmodule Xav.Frame do
     :pts
   ]
 
-  @spec new(binary(), atom(), non_neg_integer(), non_neg_integer(), integer()) :: t()
+  @doc """
+  Creates a new audio/video frame.
+  """
+  @spec new(binary(), format(), non_neg_integer(), non_neg_integer(), integer()) :: t()
   def new(data, format, width, height, pts) do
     %__MODULE__{
       type: :video,
@@ -35,7 +42,7 @@ defmodule Xav.Frame do
     }
   end
 
-  @spec new(binary(), atom(), integer(), integer()) :: t()
+  @spec new(binary(), format(), integer(), integer()) :: t()
   def new(data, format, samples, pts) do
     %__MODULE__{
       type: :audio,
@@ -47,7 +54,7 @@ defmodule Xav.Frame do
   end
 
   @doc """
-  Converts frame to Nx tensor.
+  Converts a frame to the Nx tensor.
   """
   @spec to_nx(t()) :: Nx.Tensor.t()
   def to_nx(%__MODULE__{type: :video} = frame) do

diff --git a/lib/reader.ex b/lib/reader.ex
@@ -10,7 +10,13 @@ defmodule Xav.Reader do
   Defaults to `:video`.
   * `device?` - determines whether path points to the camera. Defaults to `false`.
   """
-  @type opts :: [read: :audio | :video, device?: boolean]
+  @type opts :: [
+          read: :audio | :video,
+          device?: boolean,
+          out_format: Xav.Frame.format(),
+          out_sample_rate: integer(),
+          out_channels: integer()
+        ]
 
   @type t() :: %__MODULE__{
           reader: reference(),
@@ -49,8 +55,18 @@ defmodule Xav.Reader do
   def new(path, opts \\ []) do
     read = opts[:read] || :video
     device? = opts[:device?] || false
-
-    case Xav.Reader.NIF.new(path, to_int(device?), to_int(read)) do
+    out_format = opts[:out_format]
+    out_sample_rate = opts[:out_sample_rate] || 0
+    out_channels = opts[:out_channels] || 0
+
+    case Xav.Reader.NIF.new(
+           path,
+           to_int(device?),
+           to_int(read),
+           out_format,
+           out_sample_rate,
+           out_channels
+         ) do
       {:ok, reader, in_format, out_format, sample_rate, bit_rate, duration, codec} ->
         {:ok,
          %__MODULE__{

diff --git a/lib/reader_nif.ex b/lib/reader_nif.ex
@@ -8,7 +8,8 @@ defmodule Xav.Reader.NIF do
     :ok = :erlang.load_nif(path, 0)
   end
 
-  def new(_path, _device, _video), do: :erlang.nif_error(:undef)
+  def new(_path, _device, _video, _out_format, _out_sample_rate, _out_channels),
+    do: :erlang.nif_error(:undef)
 
   def next_frame(_reader), do: :erlang.nif_error(:undef)
 end