torchcodec 0.7.0__cp310-cp310-win_amd64.whl

torchcodec/_core/SingleStreamDecoder.h

@@ -0,0 +1,380 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <torch/types.h>
+#include <cstdint>
+#include <memory>
+#include <ostream>
+#include <string_view>
+
+#include "src/torchcodec/_core/AVIOContextHolder.h"
+#include "src/torchcodec/_core/DeviceInterface.h"
+#include "src/torchcodec/_core/FFMPEGCommon.h"
+#include "src/torchcodec/_core/Frame.h"
+#include "src/torchcodec/_core/StreamOptions.h"
+
+namespace facebook::torchcodec {
+
+// The SingleStreamDecoder class can be used to decode video frames to Tensors.
+// Note that SingleStreamDecoder is not thread-safe.
+// Do not call non-const APIs concurrently on the same object.
+class SingleStreamDecoder {
+ public:
+  // --------------------------------------------------------------------------
+  // CONSTRUCTION API
+  // --------------------------------------------------------------------------
+
+  enum class SeekMode { exact, approximate, custom_frame_mappings };
+
+  // Creates a SingleStreamDecoder from the video at videoFilePath.
+  explicit SingleStreamDecoder(
+      const std::string& videoFilePath,
+      SeekMode seekMode = SeekMode::exact);
+
+  // Creates a SingleStreamDecoder using the provided AVIOContext inside the
+  // AVIOContextHolder. The AVIOContextHolder is the base class, and the
+  // derived class will have specialized how the custom read, seek and writes
+  // work.
+  explicit SingleStreamDecoder(
+      std::unique_ptr<AVIOContextHolder> context,
+      SeekMode seekMode = SeekMode::exact);
+
+  // --------------------------------------------------------------------------
+  // VIDEO METADATA QUERY API
+  // --------------------------------------------------------------------------
+
+  // Updates the metadata of the video to accurate values obtained by scanning
+  // the contents of the video file. Also updates each StreamInfo's index, i.e.
+  // the allFrames and keyFrames vectors.
+  void scanFileAndUpdateMetadataAndIndex();
+
+  // Sorts the keyFrames and allFrames vectors in each StreamInfo by pts.
+  void sortAllFrames();
+
+  // Returns the metadata for the container.
+  ContainerMetadata getContainerMetadata() const;
+
+  // Returns the key frame indices as a tensor. The tensor is 1D and contains
+  // int64 values, where each value is the frame index for a key frame.
+  torch::Tensor getKeyFrameIndices();
+
+  // FrameMappings is used for the custom_frame_mappings seek mode to store
+  // metadata of frames in a stream. The size of all tensors in this struct must
+  // match.
+
+  // --------------------------------------------------------------------------
+  // ADDING STREAMS API
+  // --------------------------------------------------------------------------
+  struct FrameMappings {
+    // 1D tensor of int64, each value is the PTS of a frame in timebase units.
+    torch::Tensor all_frames;
+    // 1D tensor of bool, each value indicates if the corresponding frame in
+    // all_frames is a key frame.
+    torch::Tensor is_key_frame;
+    // 1D tensor of int64, each value is the duration of the corresponding frame
+    // in all_frames in timebase units.
+    torch::Tensor duration;
+  };
+
+  void addVideoStream(
+      int streamIndex,
+      const VideoStreamOptions& videoStreamOptions = VideoStreamOptions(),
+      std::optional<FrameMappings> customFrameMappings = std::nullopt);
+  void addAudioStream(
+      int streamIndex,
+      const AudioStreamOptions& audioStreamOptions = AudioStreamOptions());
+
+  // --------------------------------------------------------------------------
+  // DECODING AND SEEKING APIs
+  // --------------------------------------------------------------------------
+
+  // Places the cursor at the first frame on or after the position in seconds.
+  // Calling getNextFrame() will return the first frame at
+  // or after this position.
+  void setCursorPtsInSeconds(double seconds);
+
+  // Decodes the frame where the current cursor position is. It also advances
+  // the cursor to the next frame.
+  FrameOutput getNextFrame();
+
+  FrameOutput getFrameAtIndex(int64_t frameIndex);
+
+  // Returns frames at the given indices for a given stream as a single stacked
+  // Tensor.
+  FrameBatchOutput getFramesAtIndices(const std::vector<int64_t>& frameIndices);
+
+  // Returns frames within a given range. The range is defined by [start, stop).
+  // The values retrieved from the range are: [start, start+step,
+  // start+(2*step), start+(3*step), ..., stop). The default for step is 1.
+  FrameBatchOutput getFramesInRange(int64_t start, int64_t stop, int64_t step);
+
+  // Decodes the first frame in any added stream that is visible at a given
+  // timestamp. Frames in the video have a presentation timestamp and a
+  // duration. For example, if a frame has presentation timestamp of 5.0s and a
+  // duration of 1.0s, it will be visible in the timestamp range [5.0, 6.0).
+  // i.e. it will be returned when this function is called with seconds=5.0 or
+  // seconds=5.999, etc.
+  FrameOutput getFramePlayedAt(double seconds);
+
+  FrameBatchOutput getFramesPlayedAt(const std::vector<double>& timestamps);
+
+  // Returns frames within a given pts range. The range is defined by
+  // [startSeconds, stopSeconds) with respect to the pts values for frames. The
+  // returned frames are in pts order.
+  //
+  // Note that while stopSeconds is excluded in the half open range, this really
+  // only makes a difference when stopSeconds is exactly the pts value for a
+  // frame. Otherwise, the moment in time immediately before stopSeconds is in
+  // the range, and that time maps to the same frame as stopSeconds.
+  //
+  // The frames returned are the frames that would be played by our abstract
+  // player. Our abstract player displays frames based on pts only. It displays
+  // frame i starting at the pts for frame i, and stops at the pts for frame
+  // i+1. This model ignores a frame's reported duration.
+  //
+  // Valid values for startSeconds and stopSeconds are:
+  //
+  //   [beginStreamPtsSecondsFromContent, endStreamPtsSecondsFromContent)
+  FrameBatchOutput getFramesPlayedInRange(
+      double startSeconds,
+      double stopSeconds);
+
+  AudioFramesOutput getFramesPlayedInRangeAudio(
+      double startSeconds,
+      std::optional<double> stopSecondsOptional = std::nullopt);
+
+  class EndOfFileException : public std::runtime_error {
+   public:
+    explicit EndOfFileException(const std::string& msg)
+        : std::runtime_error(msg) {}
+  };
+
+  // --------------------------------------------------------------------------
+  // MORALLY PRIVATE APIS
+  // --------------------------------------------------------------------------
+  // These are APIs that should be private, but that are effectively exposed for
+  // practical reasons, typically for testing purposes.
+
+  // Once getFrameAtIndex supports the preAllocatedOutputTensor parameter, we
+  // can move it back to private.
+  FrameOutput getFrameAtIndexInternal(
+      int64_t frameIndex,
+      std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt);
+
+  // Exposed for _test_frame_pts_equality, which is used to test non-regression
+  // of pts resolution (64 to 32 bit floats)
+  double getPtsSecondsForFrame(int64_t frameIndex);
+
+  // Exposed for performance testing.
+  struct DecodeStats {
+    int64_t numSeeksAttempted = 0;
+    int64_t numSeeksDone = 0;
+    int64_t numSeeksSkipped = 0;
+    int64_t numPacketsRead = 0;
+    int64_t numPacketsSentToDecoder = 0;
+    int64_t numFramesReceivedByDecoder = 0;
+    int64_t numFlushes = 0;
+  };
+
+  DecodeStats getDecodeStats() const;
+  void resetDecodeStats();
+
+ private:
+  // --------------------------------------------------------------------------
+  // STREAMINFO AND ASSOCIATED STRUCTS
+  // --------------------------------------------------------------------------
+
+  struct FrameInfo {
+    int64_t pts = 0;
+
+    // The value of the nextPts default is important: the last frame's nextPts
+    // will be INT64_MAX, which ensures that the allFrames vec contains
+    // FrameInfo structs with *increasing* nextPts values. That's a necessary
+    // condition for the binary searches on those values to work properly (as
+    // typically done during pts -> index conversions).
+    // TODO: This field is unset (left to the default) for entries in the
+    // keyFrames vec!
+    int64_t nextPts = INT64_MAX;
+
+    // Note that frameIndex is ALWAYS the index into all of the frames in that
+    // stream, even when the FrameInfo is part of the key frame index. Given a
+    // FrameInfo for a key frame, the frameIndex allows us to know which frame
+    // that is in the stream.
+    int64_t frameIndex = 0;
+
+    // Indicates whether a frame is a key frame. It may appear redundant as it's
+    // only true for FrameInfos in the keyFrames index, but it is needed to
+    // correctly map frames between allFrames and keyFrames during the scan.
+    bool isKeyFrame = false;
+  };
+
+  struct StreamInfo {
+    int streamIndex = -1;
+    AVStream* stream = nullptr;
+    AVMediaType avMediaType = AVMEDIA_TYPE_UNKNOWN;
+
+    AVRational timeBase = {};
+    UniqueAVCodecContext codecContext;
+
+    // The FrameInfo indices we built when scanFileAndUpdateMetadataAndIndex was
+    // called.
+    std::vector<FrameInfo> keyFrames;
+    std::vector<FrameInfo> allFrames;
+
+    // TODO since the decoder is single-stream, these should be decoder fields,
+    // not streamInfo fields. And they should be defined right next to
+    // `cursor_`, with joint documentation.
+    int64_t lastDecodedAvFramePts = 0;
+    int64_t lastDecodedAvFrameDuration = 0;
+    VideoStreamOptions videoStreamOptions;
+    AudioStreamOptions audioStreamOptions;
+
+    // color-conversion fields. Only one of FilterGraphContext and
+    // UniqueSwsContext should be non-null.
+    UniqueSwrContext swrContext;
+  };
+
+  // --------------------------------------------------------------------------
+  // INITIALIZERS
+  // --------------------------------------------------------------------------
+
+  void initializeDecoder();
+
+  // Reads the user provided frame index and updates each StreamInfo's index,
+  // i.e. the allFrames and keyFrames vectors, and
+  // endStreamPtsSecondsFromContent
+  void readCustomFrameMappingsUpdateMetadataAndIndex(
+      int streamIndex,
+      FrameMappings customFrameMappings);
+  // --------------------------------------------------------------------------
+  // DECODING APIS AND RELATED UTILS
+  // --------------------------------------------------------------------------
+
+  void setCursor(int64_t pts);
+  void setCursor(double) = delete; // prevent calls with doubles and floats
+  bool canWeAvoidSeeking() const;
+
+  void maybeSeekToBeforeDesiredPts();
+
+  UniqueAVFrame decodeAVFrame(
+      std::function<bool(const UniqueAVFrame&)> filterFunction);
+
+  FrameOutput getNextFrameInternal(
+      std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt);
+
+  torch::Tensor maybePermuteHWC2CHW(torch::Tensor& hwcTensor);
+
+  FrameOutput convertAVFrameToFrameOutput(
+      UniqueAVFrame& avFrame,
+      std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt);
+
+  void convertAVFrameToFrameOutputOnCPU(
+      UniqueAVFrame& avFrame,
+      FrameOutput& frameOutput,
+      std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt);
+
+  void convertAudioAVFrameToFrameOutputOnCPU(
+      UniqueAVFrame& srcAVFrame,
+      FrameOutput& frameOutput);
+
+  torch::Tensor convertAVFrameToTensorUsingFilterGraph(
+      const UniqueAVFrame& avFrame);
+
+  int convertAVFrameToTensorUsingSwsScale(
+      const UniqueAVFrame& avFrame,
+      torch::Tensor& outputTensor);
+
+  std::optional<torch::Tensor> maybeFlushSwrBuffers();
+
+  // --------------------------------------------------------------------------
+  // PTS <-> INDEX CONVERSIONS
+  // --------------------------------------------------------------------------
+
+  int getKeyFrameIndexForPts(int64_t pts) const;
+
+  // Returns the key frame index of the presentation timestamp using our index.
+  // We build this index by scanning the file in
+  // scanFileAndUpdateMetadataAndIndex
+  int getKeyFrameIndexForPtsUsingScannedIndex(
+      const std::vector<SingleStreamDecoder::FrameInfo>& keyFrames,
+      int64_t pts) const;
+
+  int64_t secondsToIndexLowerBound(double seconds);
+
+  int64_t secondsToIndexUpperBound(double seconds);
+
+  int64_t getPts(int64_t frameIndex);
+
+  // --------------------------------------------------------------------------
+  // STREAM AND METADATA APIS
+  // --------------------------------------------------------------------------
+
+  void addStream(
+      int streamIndex,
+      AVMediaType mediaType,
+      const torch::Device& device = torch::kCPU,
+      std::optional<int> ffmpegThreadCount = std::nullopt);
+
+  // Returns the "best" stream index for a given media type. The "best" is
+  // determined by various heuristics in FFMPEG.
+  // See
+  // https://ffmpeg.org/doxygen/trunk/group__lavf__decoding.html#ga757780d38f482deb4d809c6c521fbcc2
+  // for more details about the heuristics.
+  // Returns the key frame index of the presentation timestamp using FFMPEG's
+  // index. Note that this index may be truncated for some files.
+  int getBestStreamIndex(AVMediaType mediaType);
+
+  std::optional<int64_t> getNumFrames(const StreamMetadata& streamMetadata);
+  double getMinSeconds(const StreamMetadata& streamMetadata);
+  std::optional<double> getMaxSeconds(const StreamMetadata& streamMetadata);
+
+  // --------------------------------------------------------------------------
+  // VALIDATION UTILS
+  // --------------------------------------------------------------------------
+
+  void validateActiveStream(
+      std::optional<AVMediaType> avMediaType = std::nullopt);
+  void validateScannedAllStreams(const std::string& msg);
+  void validateFrameIndex(
+      const StreamMetadata& streamMetadata,
+      int64_t frameIndex);
+
+  // --------------------------------------------------------------------------
+  // ATTRIBUTES
+  // --------------------------------------------------------------------------
+
+  SeekMode seekMode_;
+  ContainerMetadata containerMetadata_;
+  UniqueDecodingAVFormatContext formatContext_;
+  std::unique_ptr<DeviceInterface> deviceInterface_;
+  std::map<int, StreamInfo> streamInfos_;
+  const int NO_ACTIVE_STREAM = -2;
+  int activeStreamIndex_ = NO_ACTIVE_STREAM;
+
+  bool cursorWasJustSet_ = false;
+  // The desired position of the cursor in the stream. We send frames >= this
+  // pts to the user when they request a frame.
+  int64_t cursor_ = INT64_MIN;
+  // Stores various internal decoding stats.
+  DecodeStats decodeStats_;
+  // Stores the AVIOContext for the input buffer.
+  std::unique_ptr<AVIOContextHolder> avioContextHolder_;
+  // Whether or not we have already scanned all streams to update the metadata.
+  bool scannedAllStreams_ = false;
+  // Tracks that we've already been initialized.
+  bool initialized_ = false;
+};
+
+// Prints the SingleStreamDecoder::DecodeStats to the ostream.
+std::ostream& operator<<(
+    std::ostream& os,
+    const SingleStreamDecoder::DecodeStats& stats);
+
+SingleStreamDecoder::SeekMode seekModeFromString(std::string_view seekMode);
+
+} // namespace facebook::torchcodec

torchcodec/_core/StreamOptions.h

@@ -0,0 +1,53 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <torch/types.h>
+#include <optional>
+#include <string>
+
+namespace facebook::torchcodec {
+
+enum ColorConversionLibrary {
+  // TODO: Add an AUTO option later.
+  // Use the libavfilter library for color conversion.
+  FILTERGRAPH,
+  // Use the libswscale library for color conversion.
+  SWSCALE
+};
+
+struct VideoStreamOptions {
+  VideoStreamOptions() {}
+
+  // Number of threads we pass to FFMPEG for decoding.
+  // 0 means FFMPEG will choose the number of threads automatically to fully
+  // utilize all cores. If not set, it will be the default FFMPEG behavior for
+  // the given codec.
+  std::optional<int> ffmpegThreadCount;
+  // Currently the dimension order can be either NHWC or NCHW.
+  // H=height, W=width, C=channel.
+  std::string dimensionOrder = "NCHW";
+  // The output height and width of the frame. If not specified, the output
+  // is the same as the original video.
+  std::optional<int> width;
+  std::optional<int> height;
+  std::optional<ColorConversionLibrary> colorConversionLibrary;
+  // By default we use CPU for decoding for both C++ and python users.
+  torch::Device device = torch::kCPU;
+};
+
+struct AudioStreamOptions {
+  AudioStreamOptions() {}
+
+  // Encoding only
+  std::optional<int> bitRate;
+  // Decoding and encoding:
+  std::optional<int> numChannels;
+  std::optional<int> sampleRate;
+};
+
+} // namespace facebook::torchcodec

torchcodec/_core/ValidationUtils.cpp

@@ -0,0 +1,35 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include "src/torchcodec/_core/ValidationUtils.h"
+#include <limits>
+#include "c10/util/Exception.h"
+
+namespace facebook::torchcodec {
+
+int validateInt64ToInt(int64_t value, const std::string& parameterName) {
+  TORCH_CHECK(
+      value >= std::numeric_limits<int>::min() &&
+          value <= std::numeric_limits<int>::max(),
+      parameterName,
+      "=",
+      value,
+      " is out of range for int type.");
+
+  return static_cast<int>(value);
+}
+
+std::optional<int> validateOptionalInt64ToInt(
+    const std::optional<int64_t>& value,
+    const std::string& parameterName) {
+  if (value.has_value()) {
+    return validateInt64ToInt(value.value(), parameterName);
+  } else {
+    return std::nullopt;
+  }
+}
+
+} // namespace facebook::torchcodec

torchcodec/_core/ValidationUtils.h

@@ -0,0 +1,21 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <cstdint>
+#include <optional>
+#include <string>
+
+namespace facebook::torchcodec {
+
+int validateInt64ToInt(int64_t value, const std::string& parameterName);
+
+std::optional<int> validateOptionalInt64ToInt(
+    const std::optional<int64_t>& value,
+    const std::string& parameterName);
+
+} // namespace facebook::torchcodec

torchcodec/_core/__init__.py

@@ -0,0 +1,40 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+from ._metadata import (
+    AudioStreamMetadata,
+    ContainerMetadata,
+    get_container_metadata,
+    get_container_metadata_from_header,
+    VideoStreamMetadata,
+)
+from .ops import (
+    _add_video_stream,
+    _get_key_frame_indices,
+    _test_frame_pts_equality,
+    add_audio_stream,
+    add_video_stream,
+    create_from_bytes,
+    create_from_file,
+    create_from_file_like,
+    create_from_tensor,
+    encode_audio_to_file,
+    encode_audio_to_file_like,
+    encode_audio_to_tensor,
+    get_ffmpeg_library_versions,
+    get_frame_at_index,
+    get_frame_at_pts,
+    get_frames_at_indices,
+    get_frames_by_pts,
+    get_frames_by_pts_in_range,
+    get_frames_by_pts_in_range_audio,
+    get_frames_in_range,
+    get_json_metadata,
+    get_next_frame,
+    scan_all_streams_to_update_metadata,
+    seek_to_pts,
+)