Add stream_index parameter to Demuxer (#1659)

Author: NicolasHug

src/torchcodec/_core/Demuxer.cpp

@@ -91,13 +91,41 @@ Demuxer::Demuxer(
   select_stream(stream_index);
 }
 
+void Demuxer::validate_requested_stream(int stream_index) {
+  int num_streams = static_cast<int>(format_context_->nb_streams);
+  STD_TORCH_CHECK(
+      stream_index >= 0 && stream_index < num_streams,
+      "The stream index ",
+      stream_index,
+      " is not a valid stream. The file has ",
+      num_streams,
+      " streams, so the index must be in [0, ",
+      num_streams - 1,
+      "].");
+
+  AVMediaType media_type =
+      format_context_->streams[stream_index]->codecpar->codec_type;
+  const char* media_type_name = av_get_media_type_string(media_type);
+  STD_TORCH_CHECK(
+      media_type == AVMEDIA_TYPE_VIDEO,
+      "The stream at index ",
+      stream_index,
+      " is not a video stream, it is of type '",
+      media_type_name == nullptr ? "unknown" : media_type_name,
+      "'. Only video streams can be demuxed.");
+}
+
 void Demuxer::select_stream(std::optional<int> stream_index) {
   int status = avformat_find_stream_info(format_context_.get(), nullptr);
   STD_TORCH_CHECK(
       status >= 0,
       "Failed to find stream info: ",
       get_ffmpeg_error_string_from_error_code(status));
 
+  if (stream_index.has_value()) {
+    validate_requested_stream(*stream_index);
+  }
+
   active_stream_index_ = av_find_best_stream(
       format_context_.get(),
       AVMEDIA_TYPE_VIDEO,
@@ -107,9 +135,8 @@ void Demuxer::select_stream(std::optional<int> stream_index) {
       /*flags=*/0);
   STD_TORCH_CHECK(
       active_stream_index_ >= 0,
-      "No valid video stream found in input file (requested index ",
-      stream_index.value_or(-1),
-      ").");
+      "No valid video stream found in input file. Only video streams are "
+      "supported: audio streams cannot be demuxed.");
   stream_ = format_context_->streams[active_stream_index_];
 
   // We only need packets from the active stream, so tell FFmpeg to discard the

src/torchcodec/_core/Demuxer.h

@@ -38,9 +38,6 @@ class FORCE_PUBLIC_VISIBILITY Demuxer {
       const std::string& file_path,
       std::optional<int> stream_index = std::nullopt);
 
-  // Demuxes from the AVIOContext inside the AVIOContextHolder, which wraps an
-  // IOInterface specializing how reads and seeks work (in-memory tensor,
-  // Python file-like, ...).
   explicit Demuxer(
       std::unique_ptr<AVIOContextHolder> avio_context_holder,
       std::optional<int> stream_index = std::nullopt);
@@ -64,6 +61,7 @@ class FORCE_PUBLIC_VISIBILITY Demuxer {
   }
 
  private:
+  void validate_requested_stream(int stream_index);
   void select_stream(std::optional<int> stream_index);
 
   // Declared before format_context_ so that it outlives it: the format context

src/torchcodec/decoders/_blocks/_demuxer.py

@@ -47,8 +47,9 @@ class Demuxer:
 
         stream_index (int, optional): Specifies which stream in the video to
             demux packets from. Note that this index is absolute across all
-            media types. If left unspecified, then the :term:`best stream` is
-            used.
+            media types. It must refer to a video stream: audio streams aren't
+            supported, and requesting one raises an error. If left unspecified,
+            then the :term:`best stream` is used.
     """
 
     def __init__(

test/test_decoders.py

@@ -4255,17 +4255,26 @@ def test_cpu_fallback_is_on_cuda(self, video, expected_pix_fmt):
         assert frame.pix_fmt == expected_pix_fmt
         assert all(plane.device.type == "cuda" for plane in frame.planes)
 
-    @pytest.mark.parametrize("pix_fmt", ("pal8", "gbrpf32le"))
-    def test_planes_of_non_viewable_format(self, tmp_path, pix_fmt):
+    @pytest.mark.parametrize(
+        "pix_fmt, codec, container",
+        (
+            ("pal8", "rawvideo", "nut"),
+            # Before FFmpeg 8 the nut muxer has no rawvideo tag for the float
+            # formats and silently writes a bogus one, so the file reads back as
+            # rgb555le. EXR in mkv stores gbrpf32le properly on all versions.
+            ("gbrpf32le", "exr", "mkv"),
+        ),
+    )
+    def test_planes_of_non_viewable_format(self, tmp_path, pix_fmt, codec, container):
         # Palettised and float formats can't be handed out as views. Everything
         # *but* the planes still works, which is what lets a caller check
         # pix_fmt before reaching for them.
-        path = tmp_path / f"{pix_fmt}.nut"
+        path = tmp_path / f"{pix_fmt}.{container}"
         subprocess.run(
             [
                 "ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
                 "-f", "lavfi", "-i", "testsrc2=size=64x48:rate=10:duration=1",
-                "-c:v", "rawvideo", "-pix_fmt", pix_fmt, str(path),
+                "-c:v", codec, "-pix_fmt", pix_fmt, str(path),
             ],
             check=True,
         )  # fmt: skip
@@ -4759,6 +4768,47 @@ def test_seek_on_every_source_kind(self, make_source):
         assert got.pts_seconds == expected.pts_seconds == seconds
         assert_frames_equal(got.data, expected.data)
 
+    # ===== stream_index =====
+
+    @pytest.mark.parametrize("stream_index", (None, 0, 3))
+    def test_stream_index(self, stream_index):
+        # nasa_13013.mp4 has two video streams, 0 and 3, of different sizes,
+        # and 3 is the best one, i.e. the one used when nothing is requested.
+        demuxer = Demuxer(NASA_VIDEO.path, stream_index=stream_index)
+        decoder = PacketDecoder(demuxer)
+        converter = ColorConverter()
+        got = [
+            converter.convert(raw_frame)
+            for raw_frame in itertools.islice(
+                self._decode(decoder, self._demux(demuxer)), 10
+            )
+        ]
+
+        expected = VideoDecoder(NASA_VIDEO.path, stream_index=stream_index)[:10]
+
+        assert len(got) == len(expected) == 10
+        for got_frame, expected_data in zip(got, expected):
+            assert_frames_equal(got_frame.data, expected_data)
+
+    @pytest.mark.parametrize("stream_index", (1, 4))  # the mp4's aac streams
+    def test_audio_stream_index_raises(self, stream_index):
+        with pytest.raises(RuntimeError, match="is not a video stream.*'audio'"):
+            Demuxer(NASA_VIDEO.path, stream_index=stream_index)
+
+    def test_audio_only_file_raises(self):
+        with pytest.raises(RuntimeError, match="No valid video stream found"):
+            Demuxer(NASA_AUDIO_MP3.path)
+
+    def test_non_video_stream_index_raises(self):
+        # Stream 2 of the mp4 is a subtitle stream.
+        with pytest.raises(RuntimeError, match="is not a video stream.*'subtitle'"):
+            Demuxer(NASA_VIDEO.path, stream_index=2)
+
+    @pytest.mark.parametrize("stream_index", (-1, 6, 1000))
+    def test_invalid_stream_index_raises(self, stream_index):
+        with pytest.raises(RuntimeError, match="is not a valid stream"):
+            Demuxer(NASA_VIDEO.path, stream_index=stream_index)
+
     def test_bad_source_type_raises(self):
         with pytest.raises(TypeError, match="Unknown source type"):
             Demuxer(123)