-
Notifications
You must be signed in to change notification settings - Fork 657
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
13 changed files
with
482 additions
and
26 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,35 @@ | ||
| #include <pybind11/stl.h> | ||
| #include <torch/extension.h> | ||
| #include <torchaudio/csrc/ffmpeg/pybind/stream_reader.h> | ||
|
|
||
| namespace torchaudio { | ||
| namespace ffmpeg { | ||
| namespace { | ||
|
|
||
| PYBIND11_MODULE(_torchaudio_ffmpeg, m) { | ||
| py::class_<StreamReaderFileObj, c10::intrusive_ptr<StreamReaderFileObj>>( | ||
| m, "StreamReaderFileObj") | ||
| .def(py::init<py::object, py::object, py::object, int64_t>()) | ||
| .def("num_src_streams", &StreamReaderFileObj::num_src_streams) | ||
| .def("num_out_streams", &StreamReaderFileObj::num_out_streams) | ||
| .def( | ||
| "find_best_audio_stream", | ||
| &StreamReaderFileObj::find_best_audio_stream) | ||
| .def( | ||
| "find_best_video_stream", | ||
| &StreamReaderFileObj::find_best_video_stream) | ||
| .def("get_src_stream_info", &StreamReaderFileObj::get_src_stream_info) | ||
| .def("get_out_stream_info", &StreamReaderFileObj::get_out_stream_info) | ||
| .def("seek", &StreamReaderFileObj::seek) | ||
| .def("add_audio_stream", &StreamReaderFileObj::add_audio_stream) | ||
| .def("add_video_stream", &StreamReaderFileObj::add_video_stream) | ||
| .def("remove_stream", &StreamReaderFileObj::remove_stream) | ||
| .def("process_packet", &StreamReaderFileObj::process_packet) | ||
| .def("process_all_packets", &StreamReaderFileObj::process_all_packets) | ||
| .def("is_buffer_ready", &StreamReaderFileObj::is_buffer_ready) | ||
| .def("pop_chunks", &StreamReaderFileObj::pop_chunks); | ||
| } | ||
|
|
||
| } // namespace | ||
| } // namespace ffmpeg | ||
| } // namespace torchaudio |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,108 @@ | ||
| #include <torchaudio/csrc/ffmpeg/ffmpeg.h> | ||
| #include <torchaudio/csrc/ffmpeg/pybind/stream_reader.h> | ||
|
|
||
| namespace torchaudio { | ||
| namespace ffmpeg { | ||
| namespace { | ||
|
|
||
| static int read_function(void* opaque, uint8_t* buf, int buf_size) { | ||
| FileObj* fileobj = static_cast<FileObj*>(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<std::string>( | ||
| static_cast<py::bytes>(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<FileObj*>(opaque); | ||
| return py::cast<int64_t>(fileobj->fileobj.attr("seek")(offset, whence)); | ||
| } | ||
|
|
||
| AVIOContextPtr get_io_context(FileObj* opaque, int buffer_size) { | ||
| uint8_t* buffer = static_cast<uint8_t*>(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<void*>(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}; | ||
| } | ||
|
|
||
| c10::optional<OptionDict> convert_dict(py::object dict) { | ||
| if (dict.is_none()) | ||
| return c10::optional<OptionDict>{}; | ||
|
|
||
| c10::Dict<std::string, std::string> out; | ||
| for (std::pair<py::handle, py::handle> item : py::cast<py::dict>(dict)) { | ||
| out.insert(item.first.cast<std::string>(), item.second.cast<std::string>()); | ||
| } | ||
| return c10::optional<OptionDict>(out); | ||
| } | ||
|
|
||
| c10::optional<std::string> convert_str(py::object s) { | ||
| if (s.is_none()) | ||
| return c10::optional<std::string>{}; | ||
| return c10::optional<std::string>{ | ||
| static_cast<std::string>(py::cast<py::str>(s))}; | ||
| } | ||
|
|
||
| } // 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_, | ||
| py::object format, | ||
| py::object option, | ||
| int64_t buffer_size) | ||
| : FileObj(fileobj_, static_cast<int>(buffer_size)), | ||
| StreamReaderBinding(get_input_format_context( | ||
| "", | ||
| convert_str(format), | ||
| convert_dict(option), | ||
| pAVIO)) {} | ||
|
|
||
| } // namespace ffmpeg | ||
| } // namespace torchaudio |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,33 @@ | ||
| #pragma once | ||
| #include <torch/extension.h> | ||
| #include <torchaudio/csrc/ffmpeg/stream_reader_wrapper.h> | ||
|
|
||
| namespace torchaudio { | ||
| namespace ffmpeg { | ||
|
|
||
| // The purpose of FileObj class is so that | ||
| // FileObjStreamReader class can inherit Streamer while | ||
| // AVIOContext is initialized before AVFormat (Streamer) | ||
| struct FileObj { | ||
| py::object fileobj; | ||
| int buffer_size; | ||
| AVIOContextPtr pAVIO; | ||
| FileObj(py::object fileobj, int buffer_size); | ||
| }; | ||
|
|
||
| struct StreamReaderFileObj : public FileObj, public StreamReaderBinding { | ||
| public: | ||
| StreamReaderFileObj( | ||
| py::object fileobj, | ||
| // Note: | ||
| // Should use `py::str` or `c10::optional<std::string>` for `format`, | ||
| // and `py::dict` or `c10::optional<OptionDict>` for `option`, but | ||
| // I could not resolve TypeError related to optionality. | ||
| // So using generic `py::object`. | ||
| py::object format, | ||
| py::object option, | ||
| int64_t buffer_size); | ||
| }; | ||
|
|
||
| } // namespace ffmpeg | ||
| } // namespace torchaudio |
Oops, something went wrong.