Convert negative indices in C++, handle slices in python

Author: Daniel Flores

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -526,6 +526,12 @@ FrameOutput SingleStreamDecoder::getFrameAtIndexInternal(
   const auto& streamInfo = streamInfos_[activeStreamIndex_];
   const auto& streamMetadata =
       containerMetadata_.allStreamMetadata[activeStreamIndex_];
+
+  std::optional<int64_t> numFrames = getNumFrames(streamMetadata);
+  if (numFrames.has_value()) {
+    // If the frameIndex is negative, we convert it to a positive index
+    frameIndex = frameIndex >= 0 ? frameIndex : frameIndex + numFrames.value();
+  }
   validateFrameIndex(streamMetadata, frameIndex);
 
   int64_t pts = getPts(frameIndex);
@@ -568,8 +574,6 @@ FrameBatchOutput SingleStreamDecoder::getFramesAtIndices(
     auto indexInOutput = indicesAreSorted ? f : argsort[f];
     auto indexInVideo = frameIndices[indexInOutput];
 
-    validateFrameIndex(streamMetadata, indexInVideo);
-
     if ((f > 0) && (indexInVideo == previousIndexInVideo)) {
       // Avoid decoding the same frame twice
       auto previousIndexInOutput = indicesAreSorted ? f - 1 : argsort[f - 1];
@@ -1559,21 +1563,23 @@ void SingleStreamDecoder::validateScannedAllStreams(const std::string& msg) {
 void SingleStreamDecoder::validateFrameIndex(
     const StreamMetadata& streamMetadata,
     int64_t frameIndex) {
-  TORCH_CHECK(
-      frameIndex >= 0,
-      "Invalid frame index=" + std::to_string(frameIndex) +
-          " for streamIndex=" + std::to_string(streamMetadata.streamIndex) +
-          "; must be greater than or equal to 0");
+  if (frameIndex < 0) {
+    throw std::out_of_range(
+        "Invalid frame index=" + std::to_string(frameIndex) +
+        " for streamIndex=" + std::to_string(streamMetadata.streamIndex) +
+        "; must be greater than or equal to 0, or be a valid negative index");
+  }
 
   // Note that if we do not have the number of frames available in our metadata,
   // then we assume that the frameIndex is valid.
   std::optional<int64_t> numFrames = getNumFrames(streamMetadata);
   if (numFrames.has_value()) {
-    TORCH_CHECK(
-        frameIndex < numFrames.value(),
-        "Invalid frame index=" + std::to_string(frameIndex) +
-            " for streamIndex=" + std::to_string(streamMetadata.streamIndex) +
-            "; must be less than " + std::to_string(numFrames.value()));
+    if (frameIndex >= numFrames.value()) {
+      throw std::out_of_range(
+          "Invalid frame index=" + std::to_string(frameIndex) +
+          " for streamIndex=" + std::to_string(streamMetadata.streamIndex) +
+          "; must be less than " + std::to_string(numFrames.value()));
+    }
   }
 }
 

src/torchcodec/decoders/_video_decoder.py

@@ -195,13 +195,6 @@ def get_frame_at(self, index: int) -> Frame:
         Returns:
             Frame: The frame at the given index.
         """
-        if index < 0:
-            index += self._num_frames
-
-        if not 0 <= index < self._num_frames:
-            raise IndexError(
-                f"Index {index} is out of bounds; must be in the range [0, {self._num_frames})."
-            )
         data, pts_seconds, duration_seconds = core.get_frame_at_index(
             self._decoder, frame_index=index
         )
@@ -220,15 +213,6 @@ def get_frames_at(self, indices: list[int]) -> FrameBatch:
         Returns:
             FrameBatch: The frames at the given indices.
         """
-        for i, index in enumerate(indices):
-            index = index if index >= 0 else index + self._num_frames
-            if not 0 <= index < self._num_frames:
-                raise IndexError(
-                    f"Index {index} is out of bounds; must be in the range [0, {self._num_frames})."
-                )
-            else:
-                indices[i] = index
-
         data, pts_seconds, duration_seconds = core.get_frames_at_indices(
             self._decoder, frame_indices=indices
         )
@@ -252,8 +236,8 @@ def get_frames_in_range(self, start: int, stop: int, step: int = 1) -> FrameBatc
         Returns:
             FrameBatch: The frames within the specified range.
         """
-        start = start if start >= 0 else start + self._num_frames
-        stop = min(stop if stop >= 0 else stop + self._num_frames, self._num_frames)
+        # Adjust start / stop indices to enable indexing semantics, ex. [-10, 1000] returns the last 10 frames
+        start, stop, step = slice(start, stop, step).indices(self._num_frames)
         if not 0 <= start < self._num_frames:
             raise IndexError(
                 f"Start index {start} is out of bounds; must be in the range [0, {self._num_frames})."

test/test_decoders.py

@@ -487,10 +487,13 @@ def test_get_frame_at_tuple_unpacking(self, device):
     def test_get_frame_at_fails(self, device, seek_mode):
         decoder = VideoDecoder(NASA_VIDEO.path, device=device, seek_mode=seek_mode)
 
-        with pytest.raises(IndexError, match="out of bounds"):
+        with pytest.raises(
+            IndexError,
+            match="must be greater than or equal to 0, or be a valid negative index",
+        ):
             frame = decoder.get_frame_at(-10000)  # noqa
 
-        with pytest.raises(IndexError, match="out of bounds"):
+        with pytest.raises(IndexError, match="must be less than"):
             frame = decoder.get_frame_at(10000)  # noqa
 
     @pytest.mark.parametrize("device", cpu_and_cuda())
@@ -549,10 +552,13 @@ def test_get_frames_at(self, device, seek_mode):
     def test_get_frames_at_fails(self, device, seek_mode):
         decoder = VideoDecoder(NASA_VIDEO.path, device=device, seek_mode=seek_mode)
 
-        with pytest.raises(IndexError, match="Index -\\d+ is out of bounds"):
+        with pytest.raises(
+            IndexError,
+            match="must be greater than or equal to 0, or be a valid negative index",
+        ):
             decoder.get_frames_at([-10000])
 
-        with pytest.raises(IndexError, match="Index 390 is out of bounds"):
+        with pytest.raises(IndexError, match="Invalid frame index=390"):
             decoder.get_frames_at([390])
 
         with pytest.raises(RuntimeError, match="Expected a value of type"):
@@ -770,64 +776,56 @@ def test_get_frames_in_range(self, stream_index, device, seek_mode):
         )
 
     @pytest.mark.parametrize("device", cpu_and_cuda())
-    @pytest.mark.parametrize("stream_index", [3, None])
     @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))
-    def test_get_frames_in_range_tensor_index_semantics(
-        self, stream_index, device, seek_mode
-    ):
+    def test_get_frames_in_range_slice_indices_syntax(self, device, seek_mode):
         decoder = VideoDecoder(
             NASA_VIDEO.path,
-            stream_index=stream_index,
+            stream_index=3,
             device=device,
             seek_mode=seek_mode,
         )
-        # slices with upper bound greater than len(decoder) are supported
-        ref_frames387_389 = NASA_VIDEO.get_frame_data_by_range(
-            start=387, stop=390, stream_index=stream_index
-        ).to(device)
+
+        # high range ends get capped to num_frames
         frames387_389 = decoder.get_frames_in_range(start=387, stop=1000)
-        print(f"{frames387_389.data.shape=}")
         assert frames387_389.data.shape == torch.Size(
             [
                 3,
-                NASA_VIDEO.get_num_color_channels(stream_index=stream_index),
-                NASA_VIDEO.get_height(stream_index=stream_index),
-                NASA_VIDEO.get_width(stream_index=stream_index),
+                NASA_VIDEO.get_num_color_channels(stream_index=3),
+                NASA_VIDEO.get_height(stream_index=3),
+                NASA_VIDEO.get_width(stream_index=3),
             ]
         )
-        assert_frames_equal(ref_frames387_389, frames387_389.data)
-
-        # test that negative values in the range are supported
-        ref_frames386_389 = NASA_VIDEO.get_frame_data_by_range(
-            start=386, stop=390, stream_index=stream_index
+        ref_frame387_389 = NASA_VIDEO.get_frame_data_by_range(
+            start=387, stop=390, stream_index=3
         ).to(device)
-        frames386_389 = decoder.get_frames_in_range(start=-4, stop=1000)
-        assert frames386_389.data.shape == torch.Size(
+        assert_frames_equal(frames387_389.data, ref_frame387_389)
+
+        # negative indices are converted
+        frames387_389 = decoder.get_frames_in_range(start=-3, stop=1000)
+        assert frames387_389.data.shape == torch.Size(
             [
-                4,
-                NASA_VIDEO.get_num_color_channels(stream_index=stream_index),
-                NASA_VIDEO.get_height(stream_index=stream_index),
-                NASA_VIDEO.get_width(stream_index=stream_index),
+                3,
+                NASA_VIDEO.get_num_color_channels(stream_index=3),
+                NASA_VIDEO.get_height(stream_index=3),
+                NASA_VIDEO.get_width(stream_index=3),
             ]
         )
-        assert_frames_equal(ref_frames386_389, frames386_389.data)
-
-    @pytest.mark.parametrize("device", cpu_and_cuda())
-    @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))
-    def test_get_frames_in_range_fails(self, device, seek_mode):
-        decoder = VideoDecoder(NASA_VIDEO.path, device=device, seek_mode=seek_mode)
-
-        with pytest.raises(IndexError, match="Start index 1000 is out of bounds"):
-            decoder.get_frames_in_range(start=1000, stop=10)
+        assert_frames_equal(frames387_389.data, ref_frame387_389)
 
-        with pytest.raises(IndexError, match="Start index -\\d+ is out of bounds"):
-            decoder.get_frames_in_range(start=-1000, stop=10)
-
-        with pytest.raises(
-            IndexError,
-            match="Stop index \\(-\\d+\\) must not be less than the start index",
-        ):
-            decoder.get_frames_in_range(start=0, stop=-1000)
+        # "None" as stop is treated as end of the video
+        frames387_None = decoder.get_frames_in_range(start=-3, stop=None)
+        assert frames387_None.data.shape == torch.Size(
+            [
+                3,
+                NASA_VIDEO.get_num_color_channels(stream_index=3),
+                NASA_VIDEO.get_height(stream_index=3),
+                NASA_VIDEO.get_width(stream_index=3),
+            ]
+        )
+        reference_frame387_389 = NASA_VIDEO.get_frame_data_by_range(
+            start=387, stop=390, stream_index=3
+        ).to(device)
+        assert_frames_equal(frames387_None.data, reference_frame387_389)
 
     @pytest.mark.parametrize("device", cpu_and_cuda())
     @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))

test/test_ops.py

@@ -125,6 +125,10 @@ def test_get_frame_at_index(self, device):
             INDEX_OF_FRAME_AT_6_SECONDS
         )
         assert_frames_equal(frame6, reference_frame6.to(device))
+        # Negative indices are supported
+        frame389 = get_frame_at_index(decoder, frame_index=-1)
+        reference_frame389 = NASA_VIDEO.get_frame_data_by_index(389)
+        assert_frames_equal(frame389[0], reference_frame389.to(device))
 
     @pytest.mark.parametrize("device", cpu_and_cuda())
     def test_get_frame_with_info_at_index(self, device):
@@ -177,6 +181,32 @@ def test_get_frames_at_indices_unsorted_indices(self, device):
         with pytest.raises(AssertionError):
             assert_frames_equal(frames[0], frames[-1])
 
+    @pytest.mark.parametrize("device", cpu_and_cuda())
+    def test_get_frames_at_indices_negative_indices(self, device):
+        decoder = create_from_file(str(NASA_VIDEO.path))
+        add_video_stream(decoder, device=device)
+        frames389and387and1, *_ = get_frames_at_indices(
+            decoder, frame_indices=[-1, -3, -389]
+        )
+        reference_frame389 = NASA_VIDEO.get_frame_data_by_index(389)
+        reference_frame387 = NASA_VIDEO.get_frame_data_by_index(387)
+        reference_frame1 = NASA_VIDEO.get_frame_data_by_index(1)
+        assert_frames_equal(frames389and387and1[0], reference_frame389.to(device))
+        assert_frames_equal(frames389and387and1[1], reference_frame387.to(device))
+        assert_frames_equal(frames389and387and1[2], reference_frame1.to(device))
+
+    @pytest.mark.parametrize("device", cpu_and_cuda())
+    def test_get_frames_at_indices_fail_on_invalid_negative_indices(self, device):
+        decoder = create_from_file(str(NASA_VIDEO.path))
+        add_video_stream(decoder, device=device)
+        with pytest.raises(
+            IndexError,
+            match="must be greater than or equal to 0, or be a valid negative index",
+        ):
+            invalid_frames, *_ = get_frames_at_indices(
+                decoder, frame_indices=[-10000, -3000]
+            )
+
     @pytest.mark.parametrize("device", cpu_and_cuda())
     def test_get_frames_by_pts(self, device):
         decoder = create_from_file(str(NASA_VIDEO.path))