Refactor StreamProcessor interface (#2791)

Summary:
StreamProcessor is constructed on top of AVStream object, and attach streams defined by client code.

This commit refactor the constructor and add_stream method signature so that `add_stream`'s signature is centered around the parameters required for filter construction.

Pull Request resolved: https://github.com/pytorch/audio/pull/2791

Reviewed By: xiaohui-zhang

Differential Revision: D40667979

Pulled By: mthrok

fbshipit-source-id: 42220832f09a7895ede3cddf969d57feeb4ef7ec

torchaudio/csrc/ffmpeg/stream_reader/stream_processor.cpp

@@ -7,23 +7,22 @@ namespace ffmpeg {
 using KeyType = StreamProcessor::KeyType;
 
 StreamProcessor::StreamProcessor(
-    AVCodecParameters* codecpar,
+    AVStream* stream,
     const c10::optional<std::string>& decoder_name,
     const c10::optional<OptionDict>& decoder_option,
     const torch::Device& device)
-    : decoder(codecpar, decoder_name, decoder_option, device) {}
+    : stream(stream),
+      decoder(stream->codecpar, decoder_name, decoder_option, device) {}
 
 ////////////////////////////////////////////////////////////////////////////////
 // Configurations
 ////////////////////////////////////////////////////////////////////////////////
 KeyType StreamProcessor::add_stream(
-    AVRational input_time_base,
-    AVCodecParameters* codecpar,
     int frames_per_chunk,
     int num_chunks,
     const c10::optional<std::string>& filter_description,
     const torch::Device& device) {
-  switch (codecpar->codec_type) {
+  switch (stream->codecpar->codec_type) {
     case AVMEDIA_TYPE_AUDIO:
     case AVMEDIA_TYPE_VIDEO:
       break;
@@ -35,13 +34,12 @@ KeyType StreamProcessor::add_stream(
       std::piecewise_construct,
       std::forward_as_tuple(key),
       std::forward_as_tuple(
-          input_time_base,
-          codecpar,
+          stream->time_base,
+          stream->codecpar,
           frames_per_chunk,
           num_chunks,
           filter_description,
           device));
-  decoder_time_base = av_q2d(input_time_base);
   return key;
 }
 

torchaudio/csrc/ffmpeg/stream_reader/stream_processor.h

@@ -14,11 +14,12 @@ class StreamProcessor {
   using KeyType = int;
 
  private:
-  AVFramePtr pFrame1;
-  AVFramePtr pFrame2;
+  // Link to the corresponding stream object
+  const AVStream* stream;
 
   // Components for decoding source media
-  double decoder_time_base; // for debug
+  AVFramePtr pFrame1;
+  AVFramePtr pFrame2;
   Decoder decoder;
 
   KeyType current_key = 0;
@@ -26,7 +27,7 @@ class StreamProcessor {
 
  public:
   StreamProcessor(
-      AVCodecParameters* codecpar,
+      AVStream* stream,
       const c10::optional<std::string>& decoder_name,
       const c10::optional<OptionDict>& decoder_option,
       const torch::Device& device);
@@ -48,8 +49,6 @@ class StreamProcessor {
   // 3. Configure a buffer.
   // 4. Return filter ID.
   KeyType add_stream(
-      AVRational input_time_base,
-      AVCodecParameters* codecpar,
       int frames_per_chunk,
       int num_chunks,
       const c10::optional<std::string>& filter_description,

torchaudio/csrc/ffmpeg/stream_reader/stream_reader.cpp

@@ -251,16 +251,11 @@ void StreamReader::add_stream(
 
   if (!processors[i]) {
     processors[i] = std::make_unique<StreamProcessor>(
-        stream->codecpar, decoder, decoder_option, device);
+        stream, decoder, decoder_option, device);
   }
   stream->discard = AVDISCARD_DEFAULT;
   int key = processors[i]->add_stream(
-      stream->time_base,
-      stream->codecpar,
-      frames_per_chunk,
-      num_chunks,
-      filter_desc,
-      device);
+      frames_per_chunk, num_chunks, filter_desc, device);
   stream_indices.push_back(std::make_pair<>(i, key));
 }