add seek

pytorch · Dec 30, 2021 · 9e4fa36 · 9e4fa36
1 parent 0659c36
commit 9e4fa36
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 0 deletions.
diff --git a/torchaudio/csrc/ffmpeg/prototype.cpp b/torchaudio/csrc/ffmpeg/prototype.cpp
@@ -98,6 +98,10 @@ int64_t find_best_video_stream(S s) {
   return s->s.find_best_video_stream();
 }
 
+void seek(S s, int64_t timestamp) {
+  s->s.seek(timestamp);
+}
+
 template <typename... Args>
 std::string string_format(const std::string& format, Args... args) {
   char buffer[512];
@@ -309,6 +313,7 @@ TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
   m.def(
       "torchaudio::ffmpeg_streamer_find_best_video_stream",
       find_best_video_stream);
+  m.def("torchaudio::ffmpeg_streamer_seek", seek);
   m.def(
       "torchaudio::ffmpeg_streamer_add_basic_audio_stream",
       add_basic_audio_stream);

diff --git a/torchaudio/csrc/ffmpeg/streamer.cpp b/torchaudio/csrc/ffmpeg/streamer.cpp
@@ -133,6 +133,14 @@ bool Streamer::is_buffer_ready() const {
 ////////////////////////////////////////////////////////////////////////////////
 // Configure methods
 ////////////////////////////////////////////////////////////////////////////////
+void Streamer::seek(double timestamp) {
+  int64_t ts = static_cast<int64_t>(timestamp * AV_TIME_BASE);
+  int ret = avformat_seek_file(pFormatContext, -1, INT64_MIN, ts, INT64_MAX, 0);
+  if (ret < 0) {
+    throw std::runtime_error(std::string("Failed to seek: ") + av_err2str(ret));
+  }
+}
+
 void Streamer::add_audio_stream(
     int i,
     int frames_per_chunk,

diff --git a/torchaudio/csrc/ffmpeg/streamer.h b/torchaudio/csrc/ffmpeg/streamer.h
@@ -60,6 +60,8 @@ class Streamer {
   //////////////////////////////////////////////////////////////////////////////
   // Configure methods
   //////////////////////////////////////////////////////////////////////////////
+  void seek(double timestamp);
+
   void add_audio_stream(
       int i,
       int frames_per_chunk,