torchcodec 0.1.0

data/ext/torchcodec/Metadata.cpp

@@ -0,0 +1,124 @@
+// 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 "Metadata.h"
+#include "torch/types.h"
+
+namespace facebook::torchcodec {
+
+std::optional<double> StreamMetadata::getDurationSeconds(
+    SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      TORCH_CHECK(
+          endStreamPtsSecondsFromContent.has_value() &&
+              beginStreamPtsSecondsFromContent.has_value(),
+          "Missing beginStreamPtsSecondsFromContent or endStreamPtsSecondsFromContent");
+      return endStreamPtsSecondsFromContent.value() -
+          beginStreamPtsSecondsFromContent.value();
+    case SeekMode::approximate:
+      if (durationSecondsFromHeader.has_value()) {
+        return durationSecondsFromHeader.value();
+      }
+      if (numFramesFromHeader.has_value() && averageFpsFromHeader.has_value() &&
+          averageFpsFromHeader.value() != 0.0) {
+        return static_cast<double>(numFramesFromHeader.value()) /
+            averageFpsFromHeader.value();
+      }
+      if (durationSecondsFromContainer.has_value()) {
+        return durationSecondsFromContainer.value();
+      }
+      return std::nullopt;
+    default:
+      TORCH_CHECK(false, "Unknown SeekMode");
+  }
+}
+
+double StreamMetadata::getBeginStreamSeconds(SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      TORCH_CHECK(
+          beginStreamPtsSecondsFromContent.has_value(),
+          "Missing beginStreamPtsSecondsFromContent");
+      return beginStreamPtsSecondsFromContent.value();
+    case SeekMode::approximate:
+      if (beginStreamPtsSecondsFromContent.has_value()) {
+        return beginStreamPtsSecondsFromContent.value();
+      }
+      return 0.0;
+    default:
+      TORCH_CHECK(false, "Unknown SeekMode");
+  }
+}
+
+std::optional<double> StreamMetadata::getEndStreamSeconds(
+    SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      TORCH_CHECK(
+          endStreamPtsSecondsFromContent.has_value(),
+          "Missing endStreamPtsSecondsFromContent");
+      return endStreamPtsSecondsFromContent.value();
+    case SeekMode::approximate:
+      if (endStreamPtsSecondsFromContent.has_value()) {
+        return endStreamPtsSecondsFromContent.value();
+      }
+      return getDurationSeconds(seekMode);
+    default:
+      TORCH_CHECK(false, "Unknown SeekMode");
+  }
+}
+
+std::optional<int64_t> StreamMetadata::getNumFrames(SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      TORCH_CHECK(
+          numFramesFromContent.has_value(), "Missing numFramesFromContent");
+      return numFramesFromContent.value();
+    case SeekMode::approximate: {
+      auto durationSeconds = getDurationSeconds(seekMode);
+      if (numFramesFromHeader.has_value()) {
+        return numFramesFromHeader.value();
+      }
+      if (averageFpsFromHeader.has_value() && durationSeconds.has_value()) {
+        return static_cast<int64_t>(
+            averageFpsFromHeader.value() * durationSeconds.value());
+      }
+      return std::nullopt;
+    }
+    default:
+      TORCH_CHECK(false, "Unknown SeekMode");
+  }
+}
+
+std::optional<double> StreamMetadata::getAverageFps(SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact: {
+      auto numFrames = getNumFrames(seekMode);
+      if (numFrames.has_value() &&
+          beginStreamPtsSecondsFromContent.has_value() &&
+          endStreamPtsSecondsFromContent.has_value()) {
+        double duration = endStreamPtsSecondsFromContent.value() -
+            beginStreamPtsSecondsFromContent.value();
+        if (duration != 0.0) {
+          return static_cast<double>(numFrames.value()) / duration;
+        }
+      }
+      return averageFpsFromHeader;
+    }
+    case SeekMode::approximate:
+      return averageFpsFromHeader;
+    default:
+      TORCH_CHECK(false, "Unknown SeekMode");
+  }
+}
+
+} // namespace facebook::torchcodec

data/ext/torchcodec/Metadata.h

@@ -0,0 +1,92 @@
+// 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 <optional>
+#include <string>
+#include <vector>
+
+extern "C" {
+#include <libavcodec/avcodec.h>
+#include <libavutil/avutil.h>
+#include <libavutil/rational.h>
+}
+
+namespace facebook::torchcodec {
+
+enum class SeekMode { exact, approximate, custom_frame_mappings };
+
+struct StreamMetadata {
+  // Common (video and audio) fields derived from the AVStream.
+  int streamIndex = -1;
+
+  // See this link for what various values are available:
+  // https://ffmpeg.org/doxygen/trunk/group__lavu__misc.html#ga9a84bba4713dfced21a1a56163be1f48
+  AVMediaType mediaType = AVMEDIA_TYPE_UNKNOWN;
+
+  std::optional<AVCodecID> codecId;
+  std::optional<std::string> codecName;
+  std::optional<double> durationSecondsFromHeader;
+  std::optional<double> beginStreamSecondsFromHeader;
+  std::optional<int64_t> numFramesFromHeader;
+  std::optional<int64_t> numKeyFrames;
+  std::optional<double> averageFpsFromHeader;
+  std::optional<double> bitRate;
+
+  // Used as fallback in approximate mode when stream duration is unavailable.
+  std::optional<double> durationSecondsFromContainer;
+
+  // More accurate duration, obtained by scanning the file.
+  // These presentation timestamps are in time base.
+  std::optional<int64_t> beginStreamPtsFromContent;
+  std::optional<int64_t> endStreamPtsFromContent;
+
+  // These presentation timestamps are in seconds.
+  std::optional<double> beginStreamPtsSecondsFromContent;
+  std::optional<double> endStreamPtsSecondsFromContent;
+
+  // This can be useful for index-based seeking.
+  std::optional<int64_t> numFramesFromContent;
+
+  // Video-only fields
+  std::optional<int> width;
+  std::optional<int> height;
+  std::optional<AVRational> sampleAspectRatio;
+
+  // Audio-only fields
+  std::optional<int64_t> sampleRate;
+  std::optional<int64_t> numChannels;
+  std::optional<std::string> sampleFormat;
+
+  // Computed methods with fallback logic
+  std::optional<double> getDurationSeconds(SeekMode seekMode) const;
+  double getBeginStreamSeconds(SeekMode seekMode) const;
+  std::optional<double> getEndStreamSeconds(SeekMode seekMode) const;
+  std::optional<int64_t> getNumFrames(SeekMode seekMode) const;
+  std::optional<double> getAverageFps(SeekMode seekMode) const;
+};
+
+struct ContainerMetadata {
+  std::vector<StreamMetadata> allStreamMetadata;
+  int numAudioStreams = 0;
+  int numVideoStreams = 0;
+
+  // Note that this is the container-level duration, which is usually the max
+  // of all stream durations available in the container.
+  std::optional<double> durationSecondsFromHeader;
+
+  // Total BitRate level information at the container level in bit/s
+  std::optional<double> bitRate;
+
+  // If set, this is the index to the default audio stream.
+  std::optional<int> bestAudioStreamIndex;
+
+  // If set, this is the index to the default video stream.
+  std::optional<int> bestVideoStreamIndex;
+};
+
+} // namespace facebook::torchcodec