Merge branch 'pts_sort_and_dedup' into samplers_hack

Author: NicolasHug

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -1092,23 +1092,25 @@ VideoDecoder::BatchDecodedOutput VideoDecoder::getFramesAtIndices(
 
 VideoDecoder::BatchDecodedOutput VideoDecoder::getFramesDisplayedByTimestamps(
     int streamIndex,
-    const std::vector<double>& framePtss) {
+    const std::vector<double>& timestamps) {
   validateUserProvidedStreamIndex(streamIndex);
   validateScannedAllStreams("getFramesDisplayedByTimestamps");
 
   // The frame displayed at timestamp t and the one displayed at timestamp `t +
   // eps` are probably the same frame, with the same index. The easiest way to
   // avoid decoding that unique frame twice is to convert the input timestamps
   // to indices, and leverage the de-duplication logic of getFramesAtIndices.
+  // This means this function requires a scan.
+  // TODO: longer term, we should implement this without requiring a scan
 
   const auto& streamMetadata = containerMetadata_.streams[streamIndex];
   const auto& stream = streams_[streamIndex];
   double minSeconds = streamMetadata.minPtsSecondsFromScan.value();
   double maxSeconds = streamMetadata.maxPtsSecondsFromScan.value();
 
-  std::vector<int64_t> frameIndices(framePtss.size());
-  for (auto i = 0; i < framePtss.size(); ++i) {
-    auto framePts = framePtss[i];
+  std::vector<int64_t> frameIndices(timestamps.size());
+  for (auto i = 0; i < timestamps.size(); ++i) {
+    auto framePts = timestamps[i];
     TORCH_CHECK(
         framePts >= minSeconds && framePts < maxSeconds,
         "frame pts is " + std::to_string(framePts) + "; must be in range [" +
@@ -1119,8 +1121,8 @@ VideoDecoder::BatchDecodedOutput VideoDecoder::getFramesDisplayedByTimestamps(
         stream.allFrames.begin(),
         stream.allFrames.end() - 1,
         framePts,
-        [&stream](const FrameInfo& info, double start) {
-          return ptsToSeconds(info.nextPts, stream.timeBase) <= start;
+        [&stream](const FrameInfo& info, double framePts) {
+          return ptsToSeconds(info.nextPts, stream.timeBase) <= framePts;
         });
     int64_t frameIndex = it - stream.allFrames.begin();
     frameIndices[i] = frameIndex;

src/torchcodec/decoders/_core/VideoDecoder.h

@@ -246,7 +246,7 @@ class VideoDecoder {
 
   BatchDecodedOutput getFramesDisplayedByTimestamps(
       int streamIndex,
-      const std::vector<double>& framePtss);
+      const std::vector<double>& timestamps);
 
   // Returns frames within a given range for a given stream as a single stacked
   // Tensor. The range is defined by [start, stop). The values retrieved from

src/torchcodec/decoders/_core/VideoDecoderOps.cpp

@@ -46,7 +46,7 @@ TORCH_LIBRARY(torchcodec_ns, m) {
   m.def(
       "get_frames_by_pts_in_range(Tensor(a!) decoder, *, int stream_index, float start_seconds, float stop_seconds) -> (Tensor, Tensor, Tensor)");
   m.def(
-      "get_frames_by_pts(Tensor(a!) decoder, *, int stream_index, float[] frame_ptss) -> (Tensor, Tensor, Tensor)");
+      "get_frames_by_pts(Tensor(a!) decoder, *, int stream_index, float[] timestamps) -> (Tensor, Tensor, Tensor)");
   m.def("get_json_metadata(Tensor(a!) decoder) -> str");
   m.def("get_container_json_metadata(Tensor(a!) decoder) -> str");
   m.def(
@@ -245,11 +245,11 @@ OpsBatchDecodedOutput get_frames_in_range(
 OpsBatchDecodedOutput get_frames_by_pts(
     at::Tensor& decoder,
     int64_t stream_index,
-    at::ArrayRef<double> frame_ptss) {
+    at::ArrayRef<double> timestamps) {
   auto videoDecoder = unwrapTensorToGetDecoder(decoder);
-  std::vector<double> framePtssVec(frame_ptss.begin(), frame_ptss.end());
+  std::vector<double> timestampsVec(timestamps.begin(), timestamps.end());
   auto result =
-      videoDecoder->getFramesDisplayedByTimestamps(stream_index, framePtssVec);
+      videoDecoder->getFramesDisplayedByTimestamps(stream_index, timestampsVec);
   return makeOpsBatchDecodedOutput(result);
 }
 

src/torchcodec/decoders/_core/VideoDecoderOps.h

@@ -79,7 +79,7 @@ OpsDecodedOutput get_frame_at_pts(at::Tensor& decoder, double seconds);
 OpsBatchDecodedOutput get_frames_by_pts(
     at::Tensor& decoder,
     int64_t stream_index,
-    at::ArrayRef<double> frame_ptss);
+    at::ArrayRef<double> timestamps);
 
 // Return the frame that is visible at a given index in the video.
 OpsDecodedOutput get_frame_at_index(

src/torchcodec/decoders/_core/video_decoder_ops.py

@@ -178,7 +178,7 @@ def get_frames_by_pts_abstract(
     decoder: torch.Tensor,
     *,
     stream_index: int,
-    frame_ptss: List[float],
+    timestamps: List[float],
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     image_size = [get_ctx().new_dynamic_size() for _ in range(4)]
     return (

test/decoders/test_video_decoder_ops.py

@@ -163,16 +163,16 @@ def test_get_frames_by_pts(self):
         stream_index = 3
 
         # Note: 13.01 should give the last video frame for the NASA video
-        frame_ptss = [2, 0, 1, 0 + 1e-3, 13.01, 2 + 1e-3]
+        timestamps = [2, 0, 1, 0 + 1e-3, 13.01, 2 + 1e-3]
 
         expected_frames = [
-            get_frame_at_pts(decoder, seconds=pts)[0] for pts in frame_ptss
+            get_frame_at_pts(decoder, seconds=pts)[0] for pts in timestamps
         ]
 
         frames, *_ = get_frames_by_pts(
             decoder,
             stream_index=stream_index,
-            frame_ptss=frame_ptss,
+            timestamps=timestamps,
         )
         for frame, expected_frame in zip(frames, expected_frames):
             assert_tensor_equal(frame, expected_frame)