From b1d11719ae04daae05c09134c1851049f21d0fa7 Mon Sep 17 00:00:00 2001 From: moto <855818+mthrok@users.noreply.github.com> Date: Tue, 31 May 2022 15:21:08 -0700 Subject: [PATCH] Move FileObj to dedicated source --- torchaudio/csrc/CMakeLists.txt | 1 + .../csrc/ffmpeg/pybind/stream_reader.cpp | 70 ----------------- torchaudio/csrc/ffmpeg/pybind/stream_reader.h | 9 +-- torchaudio/csrc/ffmpeg/pybind/typedefs.cpp | 76 +++++++++++++++++++ torchaudio/csrc/ffmpeg/pybind/typedefs.h | 16 ++++ 5 files changed, 94 insertions(+), 78 deletions(-) create mode 100644 torchaudio/csrc/ffmpeg/pybind/typedefs.cpp create mode 100644 torchaudio/csrc/ffmpeg/pybind/typedefs.h diff --git a/torchaudio/csrc/CMakeLists.txt b/torchaudio/csrc/CMakeLists.txt index e73a82f39ce..086536d1aeb 100644 --- a/torchaudio/csrc/CMakeLists.txt +++ b/torchaudio/csrc/CMakeLists.txt @@ -288,6 +288,7 @@ if (BUILD_TORCHAUDIO_PYTHON_EXTENSION) if(USE_FFMPEG) set( FFMPEG_EXTENSION_SOURCES + ffmpeg/pybind/typedefs.cpp ffmpeg/pybind/pybind.cpp ffmpeg/pybind/stream_reader.cpp ) diff --git a/torchaudio/csrc/ffmpeg/pybind/stream_reader.cpp b/torchaudio/csrc/ffmpeg/pybind/stream_reader.cpp index ede150c774e..45c99e4ad31 100644 --- a/torchaudio/csrc/ffmpeg/pybind/stream_reader.cpp +++ b/torchaudio/csrc/ffmpeg/pybind/stream_reader.cpp @@ -1,77 +1,7 @@ -#include #include namespace torchaudio { namespace ffmpeg { -namespace { - -static int read_function(void* opaque, uint8_t* buf, int buf_size) { - FileObj* fileobj = static_cast(opaque); - buf_size = FFMIN(buf_size, fileobj->buffer_size); - - int num_read = 0; - while (num_read < buf_size) { - int request = buf_size - num_read; - auto chunk = static_cast( - static_cast(fileobj->fileobj.attr("read")(request))); - auto chunk_len = chunk.length(); - if (chunk_len == 0) { - break; - } - if (chunk_len > request) { - std::ostringstream message; - message - << "Requested up to " << request << " bytes but, " - << "received " << chunk_len << " bytes. " - << "The given object does not confirm to read protocol of file object."; - throw std::runtime_error(message.str()); - } - memcpy(buf, chunk.data(), chunk_len); - buf += chunk_len; - num_read += chunk_len; - } - return num_read == 0 ? AVERROR_EOF : num_read; -} - -static int64_t seek_function(void* opaque, int64_t offset, int whence) { - // We do not know the file size. - if (whence == AVSEEK_SIZE) { - return AVERROR(EIO); - } - FileObj* fileobj = static_cast(opaque); - return py::cast(fileobj->fileobj.attr("seek")(offset, whence)); -} - -AVIOContextPtr get_io_context(FileObj* opaque, int buffer_size) { - uint8_t* buffer = static_cast(av_malloc(buffer_size)); - if (!buffer) { - throw std::runtime_error("Failed to allocate buffer."); - } - - // If avio_alloc_context succeeds, then buffer will be cleaned up by - // AVIOContextPtr destructor. - // If avio_alloc_context fails, we need to clean up by ourselves. - AVIOContext* av_io_ctx = avio_alloc_context( - buffer, - buffer_size, - 0, - static_cast(opaque), - &read_function, - nullptr, - py::hasattr(opaque->fileobj, "seek") ? &seek_function : nullptr); - - if (!av_io_ctx) { - av_freep(&buffer); - throw std::runtime_error("Failed to allocate AVIO context."); - } - return AVIOContextPtr{av_io_ctx}; -} -} // namespace - -FileObj::FileObj(py::object fileobj_, int buffer_size) - : fileobj(fileobj_), - buffer_size(buffer_size), - pAVIO(get_io_context(this, buffer_size)) {} StreamReaderFileObj::StreamReaderFileObj( py::object fileobj_, diff --git a/torchaudio/csrc/ffmpeg/pybind/stream_reader.h b/torchaudio/csrc/ffmpeg/pybind/stream_reader.h index 7b12ae5c020..72a5975eab9 100644 --- a/torchaudio/csrc/ffmpeg/pybind/stream_reader.h +++ b/torchaudio/csrc/ffmpeg/pybind/stream_reader.h @@ -1,17 +1,10 @@ #pragma once -#include +#include #include namespace torchaudio { namespace ffmpeg { -struct FileObj { - py::object fileobj; - int buffer_size; - AVIOContextPtr pAVIO; - FileObj(py::object fileobj, int buffer_size); -}; - // The reason we inherit FileObj instead of making it an attribute // is so that FileObj is instantiated first. // AVIOContext must be initialized before AVFormat, and outlive AVFormat. diff --git a/torchaudio/csrc/ffmpeg/pybind/typedefs.cpp b/torchaudio/csrc/ffmpeg/pybind/typedefs.cpp new file mode 100644 index 00000000000..c358567a143 --- /dev/null +++ b/torchaudio/csrc/ffmpeg/pybind/typedefs.cpp @@ -0,0 +1,76 @@ +#include + +namespace torchaudio { +namespace ffmpeg { +namespace { + +static int read_function(void* opaque, uint8_t* buf, int buf_size) { + FileObj* fileobj = static_cast(opaque); + buf_size = FFMIN(buf_size, fileobj->buffer_size); + + int num_read = 0; + while (num_read < buf_size) { + int request = buf_size - num_read; + auto chunk = static_cast( + static_cast(fileobj->fileobj.attr("read")(request))); + auto chunk_len = chunk.length(); + if (chunk_len == 0) { + break; + } + if (chunk_len > request) { + std::ostringstream message; + message + << "Requested up to " << request << " bytes but, " + << "received " << chunk_len << " bytes. " + << "The given object does not confirm to read protocol of file object."; + throw std::runtime_error(message.str()); + } + memcpy(buf, chunk.data(), chunk_len); + buf += chunk_len; + num_read += chunk_len; + } + return num_read == 0 ? AVERROR_EOF : num_read; +} + +static int64_t seek_function(void* opaque, int64_t offset, int whence) { + // We do not know the file size. + if (whence == AVSEEK_SIZE) { + return AVERROR(EIO); + } + FileObj* fileobj = static_cast(opaque); + return py::cast(fileobj->fileobj.attr("seek")(offset, whence)); +} + +AVIOContextPtr get_io_context(FileObj* opaque, int buffer_size) { + uint8_t* buffer = static_cast(av_malloc(buffer_size)); + if (!buffer) { + throw std::runtime_error("Failed to allocate buffer."); + } + + // If avio_alloc_context succeeds, then buffer will be cleaned up by + // AVIOContextPtr destructor. + // If avio_alloc_context fails, we need to clean up by ourselves. + AVIOContext* av_io_ctx = avio_alloc_context( + buffer, + buffer_size, + 0, + static_cast(opaque), + &read_function, + nullptr, + py::hasattr(opaque->fileobj, "seek") ? &seek_function : nullptr); + + if (!av_io_ctx) { + av_freep(&buffer); + throw std::runtime_error("Failed to allocate AVIO context."); + } + return AVIOContextPtr{av_io_ctx}; +} +} // namespace + +FileObj::FileObj(py::object fileobj_, int buffer_size) + : fileobj(fileobj_), + buffer_size(buffer_size), + pAVIO(get_io_context(this, buffer_size)) {} + +} // namespace ffmpeg +} // namespace torchaudio diff --git a/torchaudio/csrc/ffmpeg/pybind/typedefs.h b/torchaudio/csrc/ffmpeg/pybind/typedefs.h new file mode 100644 index 00000000000..d1366892037 --- /dev/null +++ b/torchaudio/csrc/ffmpeg/pybind/typedefs.h @@ -0,0 +1,16 @@ +#pragma once +#include +#include + +namespace torchaudio { +namespace ffmpeg { + +struct FileObj { + py::object fileobj; + int buffer_size; + AVIOContextPtr pAVIO; + FileObj(py::object fileobj, int buffer_size); +}; + +} // namespace ffmpeg +} // namespace torchaudio