Add StreamReaderCustomIO (#3320)

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

Add StreamReaderCustomIO, which is analogous to StreamWriterCustomIO and which takes custom read/seek functions to fetch media data.

Reviewed By: hwangjeff

Differential Revision: D45482843

fbshipit-source-id: 3ccf771c0fdce153aaa7551053e9a77facedc983

torchaudio/csrc/ffmpeg/stream_reader/stream_reader.cpp

@@ -565,5 +565,47 @@ std::vector<c10::optional<Chunk>> StreamReader::pop_chunks() {
 std::vector<AVPacketPtr> StreamReader::pop_packets() {
   return packet_buffer->pop_packets();
 }
+
+//////////////////////////////////////////////////////////////////////////////
+// StreamReaderCustomIO
+//////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+namespace {
+AVIOContext* get_io_context(
+    void* opaque,
+    int buffer_size,
+    int (*read_packet)(void* opaque, uint8_t* buf, int buf_size),
+    int64_t (*seek)(void* opaque, int64_t offset, int whence)) {
+  unsigned char* buffer = static_cast<unsigned char*>(av_malloc(buffer_size));
+  TORCH_CHECK(buffer, "Failed to allocate buffer.");
+  AVIOContext* io_ctx = avio_alloc_context(
+      buffer, buffer_size, 0, opaque, read_packet, nullptr, seek);
+  if (!io_ctx) {
+    av_freep(&buffer);
+    TORCH_CHECK(false, "Failed to allocate AVIOContext.");
+  }
+  return io_ctx;
+}
+} // namespace
+
+CustomInput::CustomInput(
+    void* opaque,
+    int buffer_size,
+    int (*read_packet)(void* opaque, uint8_t* buf, int buf_size),
+    int64_t (*seek)(void* opaque, int64_t offset, int whence))
+    : io_ctx(get_io_context(opaque, buffer_size, read_packet, seek)) {}
+} // namespace detail
+
+StreamReaderCustomIO::StreamReaderCustomIO(
+    void* opaque,
+    const c10::optional<std::string>& format,
+    int buffer_size,
+    int (*read_packet)(void* opaque, uint8_t* buf, int buf_size),
+    int64_t (*seek)(void* opaque, int64_t offset, int whence),
+    const c10::optional<OptionDict>& option)
+    : CustomInput(opaque, buffer_size, read_packet, seek),
+      StreamReader(io_ctx, format, option) {}
+
 } // namespace io
 } // namespace torchaudio

torchaudio/csrc/ffmpeg/stream_reader/stream_reader.h

@@ -8,6 +8,10 @@
 namespace torchaudio {
 namespace io {
 
+//////////////////////////////////////////////////////////////////////////////
+// StreamReader
+//////////////////////////////////////////////////////////////////////////////
+
 ///
 /// Fetch and decode audio/video streams chunk by chunk.
 ///
@@ -38,13 +42,13 @@ class StreamReader {
   // to determine whether the frames should be passed to downstream.
   int64_t seek_timestamp = 0;
 
- public:
   /// @name Constructors
   ///
   ///@{
 
   /// @cond
 
+ private:
   /// Construct StreamReader from already initialized AVFormatContext.
   /// This is a low level constructor interact with FFmpeg directly.
   /// One can provide custom AVFormatContext in case the other constructor
@@ -53,36 +57,33 @@ class StreamReader {
   /// own the resources and release it at the end.
   explicit StreamReader(AVFormatContext* format_ctx);
 
-  /// @endcond
-
-  /// Construct media processor from soruce URI.
+ protected:
+  /// Concstruct media processor from custom IO.
   ///
-  /// @param src URL of source media, in the format FFmpeg can understand.
-  /// @param format Specifies format (such as mp4) or device (such as lavfi and
-  /// avfoundation)
+  /// @param io_ctx Custom IO Context.
+  /// @param format Specifies format, such as mp4.
   /// @param option Custom option passed when initializing format context
   /// (opening source).
   explicit StreamReader(
-      const std::string& src,
+      AVIOContext* io_ctx,
       const c10::optional<std::string>& format = c10::nullopt,
       const c10::optional<OptionDict>& option = c10::nullopt);
 
-  /// @cond
+  /// @endcond
 
-  /// Concstruct media processor from custom IO.
+ public:
+  /// Construct media processor from soruce URI.
   ///
-  /// @param io_ctx Custom IO Context.
-  /// @param format Specifies format, such as mp4.
+  /// @param src URL of source media, in the format FFmpeg can understand.
+  /// @param format Specifies format (such as mp4) or device (such as lavfi and
+  /// avfoundation)
   /// @param option Custom option passed when initializing format context
   /// (opening source).
-  // TODO: Move this to wrapper class
   explicit StreamReader(
-      AVIOContext* io_ctx,
+      const std::string& src,
       const c10::optional<std::string>& format = c10::nullopt,
       const c10::optional<OptionDict>& option = c10::nullopt);
 
-  /// @endcond
-
   ///@}
 
   /// @cond
@@ -339,5 +340,43 @@ class StreamReader {
   ///@}
 };
 
+//////////////////////////////////////////////////////////////////////////////
+// StreamReaderCustomIO
+//////////////////////////////////////////////////////////////////////////////
+
+/// @cond
+
+namespace detail {
+struct CustomInput {
+  AVIOContextPtr io_ctx;
+  CustomInput(
+      void* opaque,
+      int buffer_size,
+      int (*read_packet)(void* opaque, uint8_t* buf, int buf_size),
+      int64_t (*seek)(void* opaque, int64_t offset, int whence));
+};
+} // namespace detail
+
+/// @endcond
+
+/// Construct StreamReader with custom read and seek functions.
+///
+/// @param opaque Custom data used by read_packet and seek functions.
+/// @param format Specify input format.
+/// @param buffer_size The size of the intermediate buffer, which FFmpeg uses to
+/// pass data to function read_packet.
+/// @param read_packet Custom read function that is called from FFmpeg to
+/// read data from the destination.
+/// @param seek Optional seek function that is used to seek the destination.
+struct StreamReaderCustomIO : private detail::CustomInput, public StreamReader {
+  StreamReaderCustomIO(
+      void* opaque,
+      const c10::optional<std::string>& format,
+      int buffer_size,
+      int (*read_packet)(void* opaque, uint8_t* buf, int buf_size),
+      int64_t (*seek)(void* opaque, int64_t offset, int whence) = nullptr,
+      const c10::optional<OptionDict>& option = c10::nullopt);
+};
+
 } // namespace io
 } // namespace torchaudio