scaffolding

Author: NicolasHug

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -1088,6 +1088,13 @@ VideoDecoder::BatchDecodedOutput VideoDecoder::getFramesAtIndices(
   return output;
 }
 
+VideoDecoder::BatchDecodedOutput VideoDecoder::getFramesAtPtss(
+    int streamIndex,
+    const std::vector<int64_t>& framePtss,
+    const bool sortPtss) {
+        return getFramesAtIndices(streamIndex, framePtss, sortPtss);
+    }
+
 VideoDecoder::BatchDecodedOutput VideoDecoder::getFramesInRange(
     int streamIndex,
     int64_t start,

src/torchcodec/decoders/_core/VideoDecoder.h

@@ -223,6 +223,7 @@ class VideoDecoder {
   // i.e. it will be returned when this function is called with seconds=5.0 or
   // seconds=5.999, etc.
   DecodedOutput getFrameDisplayedAtTimestampNoDemux(double seconds);
+
   DecodedOutput getFrameAtIndex(
       int streamIndex,
       int64_t frameIndex,
@@ -243,6 +244,12 @@ class VideoDecoder {
       int streamIndex,
       const std::vector<int64_t>& frameIndices,
       const bool sortIndices = false);
+
+  BatchDecodedOutput getFramesAtPtss(
+      int streamIndex,
+      const std::vector<int64_t>& framePtss,
+      const bool sortPtss = false);
+
   // 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
   // the range are:

src/torchcodec/decoders/_core/VideoDecoderOps.cpp

@@ -41,6 +41,8 @@ TORCH_LIBRARY(torchcodec_ns, m) {
       "get_frame_at_index(Tensor(a!) decoder, *, int stream_index, int frame_index) -> (Tensor, Tensor, Tensor)");
   m.def(
       "get_frames_at_indices(Tensor(a!) decoder, *, int stream_index, int[] frame_indices, bool sort_indices=False) -> (Tensor, Tensor, Tensor)");
+  m.def(
+      "get_frames_at_ptss(Tensor(a!) decoder, *, int stream_index, int[] frame_ptss, bool sort_ptss=False) -> (Tensor, Tensor, Tensor)");
   m.def(
       "get_frames_in_range(Tensor(a!) decoder, *, int stream_index, int start, int stop, int? step=None) -> (Tensor, Tensor, Tensor)");
   m.def(
@@ -209,6 +211,20 @@ OpsDecodedOutput get_frame_at_pts(at::Tensor& decoder, double seconds) {
   return makeOpsDecodedOutput(result);
 }
 
+OpsBatchDecodedOutput get_frames_at_ptss(
+    at::Tensor& decoder,
+    int64_t stream_index,
+    at::IntArrayRef frame_ptss,
+    bool sort_ptss) {
+  auto videoDecoder = unwrapTensorToGetDecoder(decoder);
+  std::vector<int64_t> framePtssVec(
+      frame_ptss.begin(), frame_ptss.end());
+  auto result = videoDecoder->getFramesAtPtss(
+      stream_index, framePtssVec, sort_ptss);
+  return makeOpsBatchDecodedOutput(result);
+}
+
+
 OpsDecodedOutput get_frame_at_index(
     at::Tensor& decoder,
     int64_t stream_index,
@@ -485,6 +501,7 @@ TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) {
   m.impl("get_frame_at_pts", &get_frame_at_pts);
   m.impl("get_frame_at_index", &get_frame_at_index);
   m.impl("get_frames_at_indices", &get_frames_at_indices);
+  m.impl("get_frames_at_ptss", &get_frames_at_ptss);
   m.impl("get_frames_in_range", &get_frames_in_range);
   m.impl("get_frames_by_pts_in_range", &get_frames_by_pts_in_range);
   m.impl("_test_frame_pts_equality", &_test_frame_pts_equality);

src/torchcodec/decoders/_core/VideoDecoderOps.h

@@ -75,6 +75,13 @@ using OpsBatchDecodedOutput = std::tuple<at::Tensor, at::Tensor, at::Tensor>;
 // given timestamp T has T >= PTS and T < PTS + Duration.
 OpsDecodedOutput get_frame_at_pts(at::Tensor& decoder, double seconds);
 
+// Return the frames at given ptss for a given stream
+OpsBatchDecodedOutput get_frames_at_ptss(
+    at::Tensor& decoder,
+    int64_t stream_index,
+    at::IntArrayRef frame_ptss,
+    bool sort_ptss = false);
+
 // Return the frame that is visible at a given index in the video.
 OpsDecodedOutput get_frame_at_index(
     at::Tensor& decoder,
@@ -85,8 +92,7 @@ OpsDecodedOutput get_frame_at_index(
 // duration as tensors.
 OpsDecodedOutput get_next_frame(at::Tensor& decoder);
 
-// Return the frames at a given index for a given stream as a single stacked
-// Tensor.
+// Return the frames at given indices for a given stream
 OpsBatchDecodedOutput get_frames_at_indices(
     at::Tensor& decoder,
     int64_t stream_index,

src/torchcodec/decoders/_core/__init__.py

@@ -22,6 +22,7 @@
     get_frame_at_index,
     get_frame_at_pts,
     get_frames_at_indices,
+    get_frames_at_ptss,
     get_frames_by_pts_in_range,
     get_frames_in_range,
     get_json_metadata,

src/torchcodec/decoders/_core/video_decoder_ops.py

@@ -71,6 +71,7 @@ def load_torchcodec_extension():
 get_frame_at_pts = torch.ops.torchcodec_ns.get_frame_at_pts.default
 get_frame_at_index = torch.ops.torchcodec_ns.get_frame_at_index.default
 get_frames_at_indices = torch.ops.torchcodec_ns.get_frames_at_indices.default
+get_frames_at_ptss = torch.ops.torchcodec_ns.get_frames_at_ptss.default
 get_frames_in_range = torch.ops.torchcodec_ns.get_frames_in_range.default
 get_frames_by_pts_in_range = torch.ops.torchcodec_ns.get_frames_by_pts_in_range.default
 get_json_metadata = torch.ops.torchcodec_ns.get_json_metadata.default
@@ -171,6 +172,21 @@ def get_frame_at_pts_abstract(
         torch.empty([], dtype=torch.float),
     )
 
+@register_fake("torchcodec_ns::get_frames_at_ptss")
+def get_frames_at_pts_abstract(
+    decoder: torch.Tensor,
+    *,
+    stream_index: int,
+    frame_ptss: List[int],
+    sort_ptss: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    image_size = [get_ctx().new_dynamic_size() for _ in range(4)]
+    return (
+        torch.empty(image_size),
+        torch.empty([], dtype=torch.float),
+        torch.empty([], dtype=torch.float),
+    )
+
 
 @register_fake("torchcodec_ns::get_frame_at_index")
 def get_frame_at_index_abstract(