Fallback to DTS when PTS info is missing

Author: NicolasHug

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -26,11 +26,16 @@ int64_t secondsToClosestPts(double seconds, const AVRational& timeBase) {
       std::round(seconds * timeBase.den / timeBase.num));
 }
 
-int64_t getPtsOrDts(const UniqueAVFrame& avFrame) {
-    return avFrame->pts == INT64_MIN? avFrame->pkt_dts : avFrame->pts;
+// Some videos aren't properly encoded and do not specify pts values (for
+// packets, and thus for frames). Unset values correspond to INT64_MIN.  When
+// that happens, we fall-back to the dts value which hopefully exists and is
+// correct.
+int64_t getPtsOrDts(ReferenceAVPacket& packet) {
+  return packet->pts == INT64_MIN ? packet->dts : packet->pts;
 }
-int64_t getPtsOrDts(ReferenceAVPacket& packet){
-    return packet->pts == INT64_MIN? packet->dts : packet->pts;
+
+int64_t getPtsOrDts(const UniqueAVFrame& avFrame) {
+  return avFrame->pts == INT64_MIN ? avFrame->pkt_dts : avFrame->pts;
 }
 
 } // namespace
@@ -499,8 +504,9 @@ FrameOutput SingleStreamDecoder::getNextFrame() {
 FrameOutput SingleStreamDecoder::getNextFrameInternal(
     std::optional<torch::Tensor> preAllocatedOutputTensor) {
   validateActiveStream();
-  UniqueAVFrame avFrame = decodeAVFrame(
-      [this](const UniqueAVFrame& avFrame) { return getPtsOrDts(avFrame) >= cursor_; });
+  UniqueAVFrame avFrame = decodeAVFrame([this](const UniqueAVFrame& avFrame) {
+    return getPtsOrDts(avFrame) >= cursor_;
+  });
   return convertAVFrameToFrameOutput(avFrame, preAllocatedOutputTensor);
 }
 
@@ -620,7 +626,6 @@ FrameBatchOutput SingleStreamDecoder::getFramesInRange(
 
 FrameOutput SingleStreamDecoder::getFramePlayedAt(double seconds) {
   validateActiveStream(AVMEDIA_TYPE_VIDEO);
-  printf("seconds=%f\n", seconds);
   StreamInfo& streamInfo = streamInfos_[activeStreamIndex_];
   double frameStartTime =
       ptsToSeconds(streamInfo.lastDecodedAvFramePts, streamInfo.timeBase);
@@ -632,24 +637,15 @@ FrameOutput SingleStreamDecoder::getFramePlayedAt(double seconds) {
     // don't cache it locally, we have to rewind back.
     seconds = frameStartTime;
   }
-  printf("seconds=%f\n", seconds);
-  printf("frameStartTime=%f\n", frameStartTime);
-  printf("frameEndTime=%f\n", frameEndTime);
-  printf("TimeBase: %d %d\n", streamInfo.timeBase.num, streamInfo.timeBase.den);
-  printf("In decoding loop\n");
 
   setCursorPtsInSeconds(seconds);
   UniqueAVFrame avFrame =
       decodeAVFrame([seconds, this](const UniqueAVFrame& avFrame) {
         StreamInfo& streamInfo = streamInfos_[activeStreamIndex_];
-        double frameStartTime = ptsToSeconds(getPtsOrDts(avFrame), streamInfo.timeBase);
+        double frameStartTime =
+            ptsToSeconds(getPtsOrDts(avFrame), streamInfo.timeBase);
         double frameEndTime = ptsToSeconds(
             getPtsOrDts(avFrame) + getDuration(avFrame), streamInfo.timeBase);
-        printf("frame pts=%ld\n", avFrame->pts);
-        printf("frame pkt_dts=%ld\n", avFrame->pkt_dts);
-        printf("frameStartTime=%f\n", frameStartTime);
-        printf("frameEndTime=%f\n", frameEndTime);
-        printf("\n");
         if (frameStartTime > seconds) {
           // FFMPEG seeked past the frame we are looking for even though we
           // set max_ts to be our needed timestamp in avformat_seek_file()
@@ -1164,7 +1160,8 @@ FrameOutput SingleStreamDecoder::convertAVFrameToFrameOutput(
   FrameOutput frameOutput;
   auto& streamInfo = streamInfos_[activeStreamIndex_];
   frameOutput.ptsSeconds = ptsToSeconds(
-      getPtsOrDts(avFrame), formatContext_->streams[activeStreamIndex_]->time_base);
+      getPtsOrDts(avFrame),
+      formatContext_->streams[activeStreamIndex_]->time_base);
   frameOutput.durationSeconds = ptsToSeconds(
       getDuration(avFrame),
       formatContext_->streams[activeStreamIndex_]->time_base);

src/torchcodec/_core/_metadata.py

@@ -142,6 +142,9 @@ def average_fps(self) -> Optional[float]:
             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 / (