refactor metadata fallback logic

Author: Molly Xu

src/torchcodec/_core/CMakeLists.txt

@@ -96,6 +96,7 @@ function(make_torchcodec_libraries
         Encoder.cpp
         ValidationUtils.cpp
         Transform.cpp
+        Metadata.cpp
     )
 
     if(ENABLE_CUDA)

src/torchcodec/_core/Metadata.cpp

@@ -0,0 +1,110 @@
+// 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"
+
+namespace facebook::torchcodec {
+
+std::optional<double> StreamMetadata::getDurationSeconds(
+    SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      // In exact mode, use the scanned content value
+      if (endStreamPtsSecondsFromContent.has_value() &&
+          beginStreamPtsSecondsFromContent.has_value()) {
+        return endStreamPtsSecondsFromContent.value() -
+            beginStreamPtsSecondsFromContent.value();
+      }
+      return std::nullopt;
+    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();
+      }
+      return std::nullopt;
+  }
+  return std::nullopt;
+}
+
+double StreamMetadata::getBeginStreamSeconds(SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      if (beginStreamPtsSecondsFromContent.has_value()) {
+        return beginStreamPtsSecondsFromContent.value();
+      }
+      return 0.0;
+    case SeekMode::approximate:
+      return 0.0;
+  }
+  return 0.0;
+}
+
+std::optional<double> StreamMetadata::getEndStreamSeconds(
+    SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      if (endStreamPtsSecondsFromContent.has_value()) {
+        return endStreamPtsSecondsFromContent.value();
+      }
+      return getDurationSeconds(seekMode);
+    case SeekMode::approximate:
+      return getDurationSeconds(seekMode);
+  }
+  return std::nullopt;
+}
+
+std::optional<int64_t> StreamMetadata::getNumFrames(SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      if (numFramesFromContent.has_value()) {
+        return numFramesFromContent.value();
+      }
+      return std::nullopt;
+    case SeekMode::approximate: {
+      if (numFramesFromHeader.has_value()) {
+        return numFramesFromHeader.value();
+      }
+      if (averageFpsFromHeader.has_value() &&
+          durationSecondsFromHeader.has_value()) {
+        return static_cast<int64_t>(
+            averageFpsFromHeader.value() * durationSecondsFromHeader.value());
+      }
+      return std::nullopt;
+    }
+  }
+  return std::nullopt;
+}
+
+std::optional<double> StreamMetadata::getAverageFps(SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      if (getNumFrames(seekMode).has_value() &&
+          beginStreamPtsSecondsFromContent.has_value() &&
+          endStreamPtsSecondsFromContent.has_value() &&
+          (beginStreamPtsSecondsFromContent.value() !=
+           endStreamPtsSecondsFromContent.value())) {
+        return static_cast<double>(
+            getNumFrames(seekMode).value() /
+            (endStreamPtsSecondsFromContent.value() -
+             beginStreamPtsSecondsFromContent.value()));
+      }
+      return averageFpsFromHeader;
+    case SeekMode::approximate:
+      return averageFpsFromHeader;
+  }
+  return std::nullopt;
+}
+
+} // namespace facebook::torchcodec

src/torchcodec/_core/Metadata.h

@@ -10,6 +10,8 @@
 #include <string>
 #include <vector>
 
+#include "SeekMode.h"
+
 extern "C" {
 #include <libavcodec/avcodec.h>
 #include <libavutil/avutil.h>
@@ -52,6 +54,13 @@ struct StreamMetadata {
   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 {

src/torchcodec/_core/SeekMode.h

@@ -0,0 +1,13 @@
+// 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
+
+namespace facebook::torchcodec {
+
+enum class SeekMode { exact, approximate, custom_frame_mappings };
+
+} // namespace facebook::torchcodec

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -367,6 +367,10 @@ ContainerMetadata SingleStreamDecoder::getContainerMetadata() const {
   return containerMetadata_;
 }
 
+SeekMode SingleStreamDecoder::getSeekMode() const {
+  return seekMode_;
+}
+
 torch::Tensor SingleStreamDecoder::getKeyFrameIndices() {
   validateActiveStream(AVMEDIA_TYPE_VIDEO);
   validateScannedAllStreams("getKeyFrameIndices");

src/torchcodec/_core/SingleStreamDecoder.h

@@ -16,6 +16,7 @@
 #include "DeviceInterface.h"
 #include "FFMPEGCommon.h"
 #include "Frame.h"
+#include "SeekMode.h"
 #include "StreamOptions.h"
 #include "Transform.h"
 
@@ -30,8 +31,6 @@ class SingleStreamDecoder {
   // CONSTRUCTION API
   // --------------------------------------------------------------------------
 
-  enum class SeekMode { exact, approximate, custom_frame_mappings };
-
   // Creates a SingleStreamDecoder from the video at videoFilePath.
   explicit SingleStreamDecoder(
       const std::string& videoFilePath,
@@ -60,6 +59,9 @@ class SingleStreamDecoder {
   // Returns the metadata for the container.
   ContainerMetadata getContainerMetadata() const;
 
+  // Returns the seek mode of this decoder.
+  SeekMode getSeekMode() 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();

src/torchcodec/_core/_metadata.py

@@ -38,14 +38,22 @@ class StreamMetadata:
     stream_index: int
     """Index of the stream that this metadata refers to (int)."""
 
+    # Computed fields (computed in C++ with fallback logic)
+    duration_seconds: Optional[float]
+    """Duration of the stream in seconds. Computed in C++ with fallback logic:
+    tries to calculate from content if scan was performed, otherwise falls back
+    to header values."""
+    begin_stream_seconds: Optional[float]
+    """Beginning of the stream, in seconds. Computed in C++ with fallback logic."""
+
     def __repr__(self):
         s = self.__class__.__name__ + ":\n"
         for field in dataclasses.fields(self):
             s += f"{SPACES}{field.name}: {getattr(self, field.name)}\n"
         return s
 
 
-@dataclass
+@dataclass(repr=False)
 class VideoStreamMetadata(StreamMetadata):
     """Metadata of a single video stream."""
 
@@ -87,103 +95,19 @@ class VideoStreamMetadata(StreamMetadata):
     is the ratio between the width and height of each pixel
     (``fractions.Fraction`` or None)."""
 
-    @property
-    def duration_seconds(self) -> Optional[float]:
-        """Duration of the stream in seconds. We try to calculate the duration
-        from the actual frames if a :term:`scan` was performed. Otherwise we
-        fall back to ``duration_seconds_from_header``. If that value is also None,
-        we  instead calculate the duration from ``num_frames_from_header`` and
-        ``average_fps_from_header``.
-        """
-        if (
-            self.end_stream_seconds_from_content is not None
-            and self.begin_stream_seconds_from_content is not None
-        ):
-            return (
-                self.end_stream_seconds_from_content
-                - self.begin_stream_seconds_from_content
-            )
-        elif self.duration_seconds_from_header is not None:
-            return self.duration_seconds_from_header
-        elif (
-            self.num_frames_from_header is not None
-            and self.average_fps_from_header is not None
-        ):
-            return self.num_frames_from_header / self.average_fps_from_header
-        else:
-            return None
-
-    @property
-    def begin_stream_seconds(self) -> float:
-        """Beginning of the stream, in seconds (float). Conceptually, this
-        corresponds to the first frame's :term:`pts`. If
-        ``begin_stream_seconds_from_content`` is not None, then it is returned.
-        Otherwise, this value is 0.
-        """
-        if self.begin_stream_seconds_from_content is None:
-            return 0
-        else:
-            return self.begin_stream_seconds_from_content
-
-    @property
-    def end_stream_seconds(self) -> Optional[float]:
-        """End of the stream, in seconds (float or None).
-        Conceptually, this corresponds to last_frame.pts + last_frame.duration.
-        If ``end_stream_seconds_from_content`` is not None, then that value is
-        returned. Otherwise, returns ``duration_seconds``.
-        """
-        if self.end_stream_seconds_from_content is None:
-            return self.duration_seconds
-        else:
-            return self.end_stream_seconds_from_content
-
-    @property
-    def num_frames(self) -> Optional[int]:
-        """Number of frames in the stream (int or None).
-        This corresponds to ``num_frames_from_content`` if a :term:`scan` was made,
-        otherwise it corresponds to ``num_frames_from_header``. If that value is also
-        None, the number of frames is calculated from the duration and the average fps.
-        """
-        if self.num_frames_from_content is not None:
-            return self.num_frames_from_content
-        elif self.num_frames_from_header is not None:
-            return self.num_frames_from_header
-        elif (
-            self.average_fps_from_header is not None
-            and self.duration_seconds_from_header is not None
-        ):
-            return int(self.average_fps_from_header * self.duration_seconds_from_header)
-        else:
-            return None
-
-    @property
-    def average_fps(self) -> Optional[float]:
-        """Average fps of the stream. If a :term:`scan` was perfomed, this is
-        computed from the number of frames and the duration of the stream.
-        Otherwise we fall back to ``average_fps_from_header``.
-        """
-        if (
-            self.end_stream_seconds_from_content is None
-            or self.begin_stream_seconds_from_content is None
-            or self.num_frames is None
-            # Should never happen, but prevents ZeroDivisionError:
-            or self.end_stream_seconds_from_content
-            == self.begin_stream_seconds_from_content
-        ):
-            return self.average_fps_from_header
-        return self.num_frames / (
-            self.end_stream_seconds_from_content
-            - self.begin_stream_seconds_from_content
-        )
-
-    def __repr__(self):
-        s = super().__repr__()
-        s += f"{SPACES}duration_seconds: {self.duration_seconds}\n"
-        s += f"{SPACES}begin_stream_seconds: {self.begin_stream_seconds}\n"
-        s += f"{SPACES}end_stream_seconds: {self.end_stream_seconds}\n"
-        s += f"{SPACES}num_frames: {self.num_frames}\n"
-        s += f"{SPACES}average_fps: {self.average_fps}\n"
-        return s
+    # Computed fields (computed in C++ with fallback logic)
+    end_stream_seconds: Optional[float]
+    """End of the stream, in seconds (float or None).
+    Conceptually, this corresponds to last_frame.pts + last_frame.duration.
+    Computed in C++ with fallback logic."""
+    num_frames: Optional[int]
+    """Number of frames in the stream (int or None).
+    Computed in C++ with fallback logic: uses content if scan was performed,
+    otherwise falls back to header values or calculates from duration and fps."""
+    average_fps: Optional[float]
+    """Average fps of the stream (float or None).
+    Computed in C++ with fallback logic: if scan was performed, computes from
+    num_frames and duration, otherwise uses header value."""
 
 
 @dataclass
@@ -260,10 +184,12 @@ def get_container_metadata(decoder: torch.Tensor) -> ContainerMetadata:
         stream_dict = json.loads(_get_stream_json_metadata(decoder, stream_index))
         common_meta = dict(
             duration_seconds_from_header=stream_dict.get("durationSecondsFromHeader"),
+            duration_seconds=stream_dict.get("durationSeconds"),
             bit_rate=stream_dict.get("bitRate"),
             begin_stream_seconds_from_header=stream_dict.get(
                 "beginStreamSecondsFromHeader"
             ),
+            begin_stream_seconds=stream_dict.get("beginStreamSeconds"),
             codec=stream_dict.get("codec"),
             stream_index=stream_index,
         )
@@ -276,6 +202,9 @@ def get_container_metadata(decoder: torch.Tensor) -> ContainerMetadata:
                     end_stream_seconds_from_content=stream_dict.get(
                         "endStreamSecondsFromContent"
                     ),
+                    end_stream_seconds=stream_dict.get("endStreamSeconds"),
+                    num_frames=stream_dict.get("numFrames"),
+                    average_fps=stream_dict.get("averageFps"),
                     width=stream_dict.get("width"),
                     height=stream_dict.get("height"),
                     num_frames_from_header=stream_dict.get("numFramesFromHeader"),