Implement seeking in blocks (#1652)

Author: NicolasHug

benchmarks/bench_blocks.py

@@ -67,7 +67,7 @@ def _demux(demuxer):
 def _decode(decoder, packets):
     for packet in packets:
         yield from decoder.decode(packet)
-    yield from decoder.flush()
+    yield from decoder.drain()
 
 
 def _convert(converter, frames):

examples/decoding/blocks.py

@@ -60,7 +60,7 @@
 # ----------------
 #
 # A pipeline is just a loop. The decoder may need more than one packet before
-# it can output a frame, and it buffers a few frames that ``flush()`` returns
+# it can output a frame, and it buffers a few frames that ``drain()`` returns
 # at the end.
 #
 # ``PacketDecoder`` and ``ColorConverter`` both accept ``device="cuda"``:
@@ -76,7 +76,7 @@
 for packet in demuxer:
     for decoded_frame in packet_decoder.decode(packet):
         frames.append(color_converter.convert(decoded_frame))
-for decoded_frame in packet_decoder.flush():
+for decoded_frame in packet_decoder.drain():
     frames.append(color_converter.convert(decoded_frame))
 
 print(f"{len(frames)} frames, {frames[0].data.shape = }, "
@@ -101,7 +101,7 @@ def demux(demuxer):
 def decode(packet_decoder, packets):
     for packet in packets:
         yield from packet_decoder.decode(packet)
-    yield from packet_decoder.flush()
+    yield from packet_decoder.drain()
 
 
 def color_convert(color_converter, decoded_frames):
@@ -168,6 +168,31 @@ def demux_on_own_thread():
 # CPU while color-converting on the GPU, or feeding frames into your own
 # pre-fetching data loader.
 
+# %%
+# Seeking
+# -------
+#
+# ``Demuxer.seek()`` moves the demuxer to a timestamp. A decoder can only start
+# on a keyframe, so the seek lands on the keyframe at or before the target, and
+# the first frames that come out usually precede it: keep decoding forward and
+# drop them until you reach the timestamp you asked for.
+#
+# The seek also invalidates the frames the decoder is holding on to, so the
+# ``PacketDecoder`` must be ``reset()``.
+demuxer = Demuxer(video_path)
+packet_decoder = PacketDecoder(demuxer, device=device)
+color_converter = ColorConverter(device=device)
+
+seconds = 2.5
+demuxer.seek(seconds)
+packet_decoder.reset()
+
+frames = color_convert(color_converter, decode(packet_decoder, demux(demuxer)))
+landed_on = next(frames)
+target = next(frame for frame in frames if frame.pts_seconds >= seconds)
+print(f"asked for {seconds}s, landed on {landed_on.pts_seconds:.3f}s, "
+      f"target frame at {target.pts_seconds:.3f}s")
+
 # %%
 # Raw frames
 # ----------

src/torchcodec/_core/Demuxer.cpp

@@ -5,10 +5,32 @@
 // LICENSE file in the root directory of this source tree.
 
 #include "Demuxer.h"
+
+#include <sstream>
+
 #include "StableABICompat.h"
 
 namespace facebook::torchcodec {
 
+// FFmpeg reports "this seek cannot be performed" as a bare -1, i.e. EPERM,
+// which renders as the very misleading "Operation not permitted". It covers
+// both a target that the demuxer can't reach and a demuxer with no seeking
+// support whatsoever.
+std::string get_seek_error_message(
+    const AVFormatContext* format_context,
+    int64_t desired_pts,
+    int status) {
+  std::stringstream ss;
+  ss << "Could not seek file to pts=" << desired_pts << ": "
+     << get_ffmpeg_error_string_from_error_code(status) << ".";
+  if (status == AVERROR(EPERM)) {
+    ss << " This is either because that timestamp is out of range, or because"
+       << " the '" << format_context->iformat->name << "' format does not"
+       << " support seeking.";
+  }
+  return ss.str();
+}
+
 int read_next_packet(
     AVFormatContext* format_context,
     int active_stream_index,
@@ -68,6 +90,21 @@ Demuxer::Demuxer(
   }
 }
 
+void Demuxer::seek(double seconds) {
+  int64_t desired_pts = seconds_to_closest_pts(seconds, stream_->time_base);
+
+  int status = avformat_seek_file(
+      format_context_.get(),
+      active_stream_index_,
+      INT64_MIN,
+      desired_pts,
+      desired_pts,
+      0);
+  STD_TORCH_CHECK(
+      status >= 0,
+      get_seek_error_message(format_context_.get(), desired_pts, status));
+}
+
 UniqueAVPacket Demuxer::next_packet() {
   // TODO_API_BREAKDOWN CC P2: Not a fan of the ReferenceAVPacket / AutoAVPacket
   // / UniqueAVPacket dance here. Can we simplify?

src/torchcodec/_core/Demuxer.h

@@ -23,6 +23,11 @@ int read_next_packet(
     int active_stream_index,
     ReferenceAVPacket& packet);
 
+std::string get_seek_error_message(
+    const AVFormatContext* format_context,
+    int64_t desired_pts,
+    int status);
+
 // Demux building block: owns an AVFormatContext, selects one video stream, and
 // yields its (compressed) packets. Does no decoding. Not thread-safe.
 class FORCE_PUBLIC_VISIBILITY Demuxer {
@@ -35,6 +40,8 @@ class FORCE_PUBLIC_VISIBILITY Demuxer {
   // packet, or a null packet at end of stream.
   UniqueAVPacket next_packet();
 
+  void seek(double seconds);
+
   AVStream* active_stream() const {
     return stream_;
   }

src/torchcodec/_core/PacketDecoder.cpp

@@ -65,6 +65,8 @@ PacketDecoder::PacketDecoder(
 
   AVStream* stream = demuxer.active_stream();
   time_base_ = stream->time_base;
+  is_mpeg_ps_ =
+      std::string_view(demuxer.format_context()->iformat->name) == "mpeg";
   if (const int32_t* matrix = get_display_matrix_from_stream(stream)) {
     display_matrix_.emplace();
     std::copy(
@@ -100,13 +102,33 @@ int PacketDecoder::send_packet(AVPacket* packet) {
   ReferenceAVPacket ref(auto_packet);
   int status = av_packet_ref(ref.get(), packet);
   STD_TORCH_CHECK(status >= AVSUCCESS, "av_packet_ref failed");
-  return device_interface_->send_packet(ref);
+
+  status = device_interface_->send_packet(ref);
+
+  if (status == AVERROR_INVALIDDATA && packet_data_may_be_misaligned_) {
+    // Seeking in an MPEG program stream lands on a container-level byte offset,
+    // so the parser resumes mid-frame and the packets it rebuilds are garbage
+    // until it resyncs. Report those as consumed rather than as a corrupt file:
+    // dropping them is exactly what resyncing means.
+    return AVSUCCESS;
+  }
+  if (status >= AVSUCCESS) {
+    // The decoder accepted a packet, so we're aligned again: from now on
+    // invalid data means the file is corrupt, and we want to report it.
+    packet_data_may_be_misaligned_ = false;
+  }
+  return status;
 }
 
 int PacketDecoder::send_eof() {
   return device_interface_->send_eof_packet();
 }
 
+void PacketDecoder::reset() {
+  device_interface_->flush();
+  packet_data_may_be_misaligned_ = is_mpeg_ps_;
+}
+
 int PacketDecoder::receive_frame(UniqueAVFrame& av_frame) {
   int status = device_interface_->receive_frame(av_frame);
   if (status == AVSUCCESS) {

src/torchcodec/_core/PacketDecoder.h

@@ -45,6 +45,13 @@ class FORCE_PUBLIC_VISIBILITY PacketDecoder {
   // Pull one frame. Returns AVSUCCESS with `av_frame` filled, AVERROR(EAGAIN)
   // if more input is needed, AVERROR_EOF at end, or a negative error code.
   int receive_frame(UniqueAVFrame& av_frame);
+  // Drop the codec's buffered state (reference frames, in-flight frames) and
+  // start over. Needed after the demuxer seeked, and after send_eof(), which
+  // otherwise leaves the codec permanently in its drained state.
+  // This is called 'reset()' and not 'flush()', because this is publicly
+  // exposed flush() is slightly ambiguous and could mean 'flush the frames out
+  // of the decoder' rather than meaning 'flush the decoder internal state'.
+  void reset();
 
   std::optional<torch::stable::Tensor> get_frame_storage(
       const AVFrame& av_frame) const {
@@ -69,6 +76,10 @@ class FORCE_PUBLIC_VISIBILITY PacketDecoder {
   // value: we're only handed the Demuxer at construction and it may well be
   // gone by the time we decode.
   std::optional<std::array<int32_t, 9>> display_matrix_;
+  // The MPEG-PS demuxer doesn't return proper packets just after a seek, so the
+  // first ones we're fed may not be decodable. See send_packet().
+  bool is_mpeg_ps_ = false;
+  bool packet_data_may_be_misaligned_ = false;
 };
 
 // A decoded frame's own samples, before any color conversion.

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -25,29 +25,6 @@ extern "C" {
 
 namespace facebook::torchcodec {
 
-namespace {
-
-// FFmpeg reports "this seek cannot be performed" as a bare -1, i.e. EPERM,
-// which renders as the very misleading "Operation not permitted". It covers
-// both a target that the demuxer can't reach and a demuxer with no seeking
-// support whatsoever.
-std::string get_seek_error_message(
-    const AVFormatContext* format_context,
-    int64_t desired_pts,
-    int status) {
-  std::stringstream ss;
-  ss << "Could not seek file to pts=" << desired_pts << ": "
-     << get_ffmpeg_error_string_from_error_code(status) << ".";
-  if (status == AVERROR(EPERM)) {
-    ss << " This is either because that timestamp is out of range, or because"
-       << " the '" << format_context->iformat->name << "' format does not"
-       << " support seeking.";
-  }
-  return ss.str();
-}
-
-} // namespace
-
 // --------------------------------------------------------------------------
 // CONSTRUCTORS, INITIALIZATION, DESTRUCTORS
 // --------------------------------------------------------------------------

src/torchcodec/_core/_ffmpeg_op_names.py

@@ -31,9 +31,11 @@
         "get_json_metadata",
         "_blocks_create_demuxer",
         "_blocks_demuxer_next_packet",
+        "_blocks_demuxer_seek",
         "_blocks_create_packet_decoder",
         "_blocks_packet_decoder_send_packet",
         "_blocks_packet_decoder_send_eof",
+        "_blocks_packet_decoder_reset",
         "_blocks_packet_decoder_receive_frame",
         "_blocks_create_color_converter",
         "_blocks_convert_frame",

src/torchcodec/_core/_ffmpeg_ops.py

@@ -96,6 +96,7 @@ def add_video_stream(
 _blocks_demuxer_next_packet = (
     torch.ops.torchcodec_ns._blocks_demuxer_next_packet.default
 )
+_blocks_demuxer_seek = torch.ops.torchcodec_ns._blocks_demuxer_seek.default
 _blocks_create_packet_decoder = (
     torch.ops.torchcodec_ns._blocks_create_packet_decoder.default
 )
@@ -105,6 +106,9 @@ def add_video_stream(
 _blocks_packet_decoder_send_eof = (
     torch.ops.torchcodec_ns._blocks_packet_decoder_send_eof.default
 )
+_blocks_packet_decoder_reset = (
+    torch.ops.torchcodec_ns._blocks_packet_decoder_reset.default
+)
 _blocks_packet_decoder_receive_frame = (
     torch.ops.torchcodec_ns._blocks_packet_decoder_receive_frame.default
 )

src/torchcodec/_core/custom_ops.cpp

@@ -76,11 +76,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(torchcodec_ns, m) {
   m.def(
       "_blocks_create_demuxer(str filename, int? stream_index=None) -> Tensor");
   m.def("_blocks_demuxer_next_packet(Tensor(a!) demuxer) -> (Tensor, bool)");
+  m.def("_blocks_demuxer_seek(Tensor(a!) demuxer, float seconds) -> ()");
   m.def(
       "_blocks_create_packet_decoder(Tensor demuxer, *, int? num_threads=None, str device=\"cpu\") -> Tensor");
   m.def(
       "_blocks_packet_decoder_send_packet(Tensor(a!) decoder, Tensor packet) -> int");
   m.def("_blocks_packet_decoder_send_eof(Tensor(a!) decoder) -> int");
+  m.def("_blocks_packet_decoder_reset(Tensor(a!) decoder) -> ()");
   m.def(
       "_blocks_packet_decoder_receive_frame(Tensor(a!) decoder) -> (Tensor, int, float, float, str, Tensor)");
   m.def(
@@ -815,6 +817,10 @@ OpsPacketOutput _blocks_demuxer_next_packet(torch::stable::Tensor& demuxer) {
   return std::make_tuple(wrap_pointer_to_tensor(std::move(packet)), false);
 }
 
+void _blocks_demuxer_seek(torch::stable::Tensor& demuxer, double seconds) {
+  unwrap_tensor_to_pointer<Demuxer>(demuxer)->seek(seconds);
+}
+
 torch::stable::Tensor _blocks_create_packet_decoder(
     torch::stable::Tensor& demuxer,
     std::optional<int64_t> num_threads,
@@ -844,6 +850,10 @@ int64_t _blocks_packet_decoder_send_eof(torch::stable::Tensor& decoder) {
   return static_cast<int64_t>(decoder_ptr->send_eof());
 }
 
+void _blocks_packet_decoder_reset(torch::stable::Tensor& decoder) {
+  unwrap_tensor_to_pointer<PacketDecoder>(decoder)->reset();
+}
+
 // TODO_API_BREAKDOWN CC P1: I hate this.
 std::string device_to_string(const StableDevice& device) {
   std::string name = device_type_name(device.type());
@@ -1498,6 +1508,7 @@ STABLE_TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) {
   m.impl("get_frames_by_pts", TORCH_BOX(&get_frames_by_pts));
   m.impl(
       "_blocks_demuxer_next_packet", TORCH_BOX(&_blocks_demuxer_next_packet));
+  m.impl("_blocks_demuxer_seek", TORCH_BOX(&_blocks_demuxer_seek));
   m.impl(
       "_blocks_create_packet_decoder",
       TORCH_BOX(&_blocks_create_packet_decoder));
@@ -1507,6 +1518,8 @@ STABLE_TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) {
   m.impl(
       "_blocks_packet_decoder_send_eof",
       TORCH_BOX(&_blocks_packet_decoder_send_eof));
+  m.impl(
+      "_blocks_packet_decoder_reset", TORCH_BOX(&_blocks_packet_decoder_reset));
   m.impl(
       "_blocks_packet_decoder_receive_frame",
       TORCH_BOX(&_blocks_packet_decoder_receive_frame));